+#! /usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2020-2023 Alibaba Group Holding Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import contextlib
+import logging
+import os
+import sys
+from typing import Any
+from typing import Dict
+from typing import Generator
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
+from .version import __version__
+
+__doc__ = """
+Vineyard - an in-memory immutable data manager. (Project under CNCF)
+====================================================================
+
+Vineyard (v6d) is an in-memory immutable data manager that provides
+out-of-the-box high-level abstraction and zero-copy in-memory
+sharing for distributed data in big data tasks, such as graph analytics
+(e.g., GraphScope), numerical computing (e.g., Mars), and machine learning.
+"""
+
+# pylint: disable=import-outside-toplevel,wrong-import-position
+
+logger = logging.getLogger('vineyard')
+
+
+@contextlib.contextmanager
+def envvars(
+ key: Union[str, Dict[str, str]],
+ value: Union[str, None] = None,
+ append: bool = False,
+) -> Generator[os._Environ, Any, Any]:
+ """Create a context with specified environment variables set.
+
+ It is useful for setting the :code`VINEYARD_IPC_SOCKET` environment
+ variable to obtain a proper default vineyard client.
+
+ This context macro can be used as
+
+ .. code:: python
+
+ with environment('KEY'):
+ # env :code:`KEY` will be set to None.
+
+ with environment('KEY', 'value'):
+ # env :code:`KEY` will be set as :code:`value`.
+
+ with environment({'KEY1': None, 'KEY2': 'value2'}):
+ # env :code:`KEY1` will be set as None and :code:`KEY2` will
+ # be set as :code:`value2`.
+ """
+ if isinstance(key, str):
+ if value is None:
+ items = dict()
+ else:
+ items: Dict[str, str] = {key: value}
+ else:
+ items: Dict[str, str] = key
+ original_items = dict()
+ for k, v in items.items():
+ original_items[k] = os.environ.get(k, None)
+ if append and original_items[k] is not None:
+ os.environ[k] = original_items[k] + ':' + v
+ else:
+ os.environ[k] = v
+
+ yield os.environ
+
+ for k, v in original_items.items():
+ if v is not None:
+ os.environ[k] = v
+ else:
+ del os.environ[k]
+
+
+def _init_global_context():
+ import os as _dl_flags # pylint: disable=reimported
+
+ if sys.platform == 'linux':
+ registry = os.path.join(
+ os.path.dirname(__file__), 'libvineyard_internal_registry.so'
+ )
+ elif sys.platform == 'darwin':
+ registry = os.path.join(
+ os.path.dirname(__file__), 'libvineyard_internal_registry.dylib'
+ )
+ else:
+ raise RuntimeError("Unsupported platform: %s" % sys.platform)
+
+ ctx = {'__VINEYARD_INTERNAL_REGISTRY': registry}
+
+ if os.environ.get('VINEYARD_DEVELOP', None) is None:
+ with envvars(ctx): # n.b., no append
+ from . import _C
+ return
+
+ if not hasattr(_dl_flags, 'RTLD_GLOBAL') or not hasattr(_dl_flags, 'RTLD_LAZY'):
+ try:
+ # next try if DLFCN exists
+ import DLFCN as _dl_flags # noqa: N811
+ except ImportError:
+ _dl_flags = None
+
+ if _dl_flags is not None:
+ old_flags = sys.getdlopenflags()
+
+ # import the extension module
+ sys.setdlopenflags(_dl_flags.RTLD_GLOBAL | _dl_flags.RTLD_LAZY)
+ with envvars(ctx): # n.b., no append
+ from . import _C # noqa: F811
+ # restore
+ sys.setdlopenflags(old_flags)
+
+
+_init_global_context()
+del _init_global_context
+
+
+from . import core
+from . import csi
+from . import data
+from . import deploy
+from . import io
+from . import launcher
+from . import shared_memory
+from ._C import ArrowErrorException
+from ._C import AssertionFailedException
+from ._C import Blob
+from ._C import BlobBuilder
+from ._C import ConnectionErrorException
+from ._C import ConnectionFailedException
+from ._C import EndOfFileException
+from ._C import EtcdErrorException
+from ._C import InstanceStatus
+from ._C import InvalidException
+from ._C import InvalidStreamStateException
+from ._C import IOErrorException
+from ._C import IPCClient
+from ._C import KeyErrorException
+from ._C import MetaTreeInvalidException
+from ._C import MetaTreeLinkInvalidException
+from ._C import MetaTreeNameInvalidException
+from ._C import MetaTreeNameNotExistsException
+from ._C import MetaTreeSubtreeNotExistsException
+from ._C import MetaTreeTypeInvalidException
+from ._C import MetaTreeTypeNotExistsException
+from ._C import NotEnoughMemoryException
+from ._C import NotImplementedException
+from ._C import Object
+from ._C import ObjectBuilder
+from ._C import ObjectExistsException
+from ._C import ObjectID
+from ._C import ObjectMeta
+from ._C import ObjectName
+from ._C import ObjectNotExistsException
+from ._C import ObjectNotSealedException
+from ._C import ObjectSealedException
+from ._C import RemoteBlob
+from ._C import RemoteBlobBuilder
+from ._C import RPCClient
+from ._C import StreamDrainedException
+from ._C import StreamFailedException
+from ._C import TypeErrorException
+from ._C import UnknownErrorException
+from ._C import UserInputErrorException
+from ._C import VineyardServerNotReadyException
+from ._C import _connect
+from ._C import memory_copy
+from .core import Client
+from .core import builder_context
+from .core import default_builder_context
+from .core import default_driver_context
+from .core import default_resolver_context
+from .core import driver_context
+from .core import resolver_context
+from .core.builder import BuilderContext
+from .core.resolver import ResolverContext
+from .data import register_builtin_types
+from .data.graph import Graph
+from .deploy.local import get_current_socket
+from .deploy.local import shutdown
+from .deploy.local import try_init
+
+
+def _init_vineyard_modules(): # noqa: C901
+ """Resolve registered vineyard modules in the following order:
+
+ * /etc/vineyard/config.py
+ * {sys.prefix}/etc/vineyard/config.py
+ * /usr/share/vineyard/01-xxx.py
+ * /usr/local/share/vineyard/01-xxx.py
+ * {sys.prefix}/share/vineyard/02-xxxx.py
+ * $HOME/.vineyard/03-xxxxx.py
+
+ Then import packages like vineyard.drivers.*:
+
+ * vineyard.drivers.io
+ """
+
+ import glob
+ import importlib.util
+ import pkgutil
+ import site
+ import sysconfig
+
+ def _import_module_from_file(filepath):
+ filepath = os.path.expanduser(os.path.expandvars(filepath))
+ if os.path.exists(filepath):
+ try:
+ spec = importlib.util.spec_from_file_location(
+ "vineyard._contrib", filepath
+ )
+ mod = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(mod)
+ except Exception: # pylint: disable=broad-except
+ logger.debug("Failed to load %s", filepath, exc_info=True)
+
+ def _import_module_from_qualified_name(module):
+ try:
+ importlib.import_module(module)
+ except Exception: # pylint: disable=broad-except
+ logger.debug('Failed to load module %s', module, exc_info=True)
+
+ _import_module_from_file('/etc/vineyard/config.py')
+ _import_module_from_file(os.path.join(sys.prefix, '/etc/vineyard/config.py'))
+ for filepath in glob.glob('/usr/share/vineyard/*-*.py'):
+ _import_module_from_file(filepath)
+ for filepath in glob.glob('/usr/local/share/vineyard/*-*.py'):
+ _import_module_from_file(filepath)
+ for filepath in glob.glob(os.path.join(sys.prefix, '/share/vineyard/*-*.py')):
+ _import_module_from_file(filepath)
+ for filepath in glob.glob(os.path.expanduser('$HOME/.vineyard/*-*.py')):
+ _import_module_from_file(filepath)
+
+ package_sites = set()
+ pkg_sites = site.getsitepackages()
+ if not isinstance(pkg_sites, (list, tuple)):
+ pkg_sites = [pkg_sites]
+ package_sites.update(pkg_sites)
+ pkg_sites = site.getusersitepackages()
+ if not isinstance(pkg_sites, (list, tuple)):
+ pkg_sites = [pkg_sites]
+ package_sites.update(pkg_sites)
+
+ paths = sysconfig.get_paths()
+ if 'purelib' in paths:
+ package_sites.add(paths['purelib'])
+ if 'platlib' in paths:
+ package_sites.add(paths['purelib'])
+
+ # add relative path
+ package_sites.add(os.path.join(os.path.dirname(__file__), '..'))
+
+ # dedup
+ deduped = set()
+ for pkg_site in package_sites:
+ deduped.add(os.path.abspath(pkg_site))
+
+ for pkg_site in deduped:
+ for _, mod, _ in pkgutil.iter_modules(
+ [os.path.join(pkg_site, 'vineyard', 'drivers')]
+ ):
+ _import_module_from_qualified_name('vineyard.drivers.%s' % mod)
+
+
+try:
+ _init_vineyard_modules()
+except Exception: # pylint: disable=broad-except
+ pass
+del _init_vineyard_modules
+
+
+[docs]def connect(*args, **kwargs):
+ """
+ Connect to vineyard by specified UNIX-domain socket or TCP endpoint.
+
+ If no arguments are provided and failed to resolve both the environment
+ variables :code:`VINEYARD_IPC_SOCKET`, :code:`VINEYARD_RPC_ENDPOINT`,
+ :code:`VINEYARD_CONFIG`, and the default configuration file
+ :code:`/var/run/vineyard-config.yaml` and
+ :code:`/var/run/vineyard/vineyard-config.yaml`, it will launch a standalone
+ vineyardd server in the background and then connect to it.
+
+ The `connect()` method has various overloading:
+
+ .. function:: connect(socket: str,
+ username: str = None,
+ password: str = None) -> IPCClient
+ :noindex:
+
+ Connect to vineyard via UNIX domain socket for IPC service:
+
+ .. code:: python
+
+ client = vineyard.connect('/var/run/vineyard.sock')
+
+ Parameters:
+ socket: str
+ UNIX domain socket path to setup an IPC connection.
+ username: str
+ Username to login, default to None.
+ password: str
+ Password to login, default to None.
+
+ Returns:
+ IPCClient: The connected IPC client.
+
+ .. function:: connect(host: str,
+ port: int or str,
+ username: str = None,
+ password: str = None) -> RPCClient
+ :noindex:
+
+ Connect to vineyard via TCP socket.
+
+ Parameters:
+ host: str
+ Hostname to connect to.
+ port: int or str
+ The TCP that listened by vineyard TCP service.
+ username: str
+ Username to login, default to None.
+ password: str
+ Password to login, default to None.
+
+ Returns:
+ RPCClient: The connected RPC client.
+
+ .. function:: connect(endpoint: (str, int or str),
+ username: str = None,
+ password: str = None) -> RPCClient
+ :noindex:
+
+ Connect to vineyard via TCP socket.
+
+ Parameters:
+ endpoint: tuple(str, int or str)
+ Endpoint to connect to. The parameter is a tuple, in which the first
+ element is the host, and the second parameter is the port (can be int
+ a str).
+ username: str
+ Username to login, default to None.
+ password: str
+ Password to login, default to None.
+
+ Returns:
+ RPCClient: The connected RPC client.
+
+ .. function:: connect(username: str = None,
+ password: str = None) -> IPCClient or RPCClient
+ :noindex:
+
+ Connect to vineyard via UNIX domain socket or TCP endpoint. This method normally
+ usually no arguments, and will first tries to resolve IPC socket from the
+ environment variable `VINEYARD_IPC_SOCKET` and connect to it. If it fails to
+ establish a connection with vineyard server, the method will tries to resolve
+ RPC endpoint from the environment variable `VINEYARD_RPC_ENDPOINT`. If both
+ tries are failed, this method will try to resolve the configuration file that
+ contains IPC socket and RPC endpoint from the environment variable
+ `VINEYARD_CONFIG`, and then connect to the vineyard server with the
+ resolved configuration.
+
+ If above all are failed, this method will raise a :class:`ConnectionFailed`
+ exception.
+
+ In rare cases, user may be not sure about if the IPC socket or RPC endpoint
+ is available, i.e., the variable might be :code:`None`. In such cases this
+ method can accept a `None` as arguments, and do resolution as described above.
+
+ Parameters:
+ username: str
+ Username to login, default to None.
+ password: str
+ Password to login, default to None.
+
+ Raises:
+ ConnectionFailed
+ """
+ if (
+ not args
+ and not kwargs
+ and 'VINEYARD_IPC_SOCKET' not in os.environ
+ and 'VINEYARD_RPC_ENDPOINT' not in os.environ
+ and 'VINEYARD_CONFIG' not in os.environ
+ and not any(
+ os.path.exists(path)
+ for path in [
+ "/var/run/vineyard-config.yaml",
+ "/var/run/vineyard/vineyard-config.yaml",
+ ]
+ )
+ ):
+ logger.info(
+ 'No vineyard socket or endpoint is specified, '
+ 'try to launch a standalone one.'
+ )
+ try_init()
+ return Client(*args, **kwargs)
+
+
+def put(
+ value: Any,
+ builder: Optional[BuilderContext] = None,
+ persist: bool = False,
+ name: Optional[str] = None,
+ **kwargs,
+):
+ """
+ Connect the vineyard server by the following Environment Variables:
+
+ VINEYARD_IPC_SOCKET:
+ UNIX domain socket path to setup an IPC connection.
+ E.g. /var/run/vineyard.sock
+ VINEYARD_RPC_ENDPOINT:
+ TCP endpoint to setup an RPC connection.
+ E.g. 127.0.0.1:9600
+ VINEYARD_CONFIG:
+ Either be a path to a YAML configuration file or a path to a
+ directory containing the default config file `vineyard-config.yaml`.
+
+ The configuration file should be like:
+
+ .. code:: yaml
+
+ Vineyard:
+ IPCSocket: '/path/to/vineyard.sock'
+ RPCEndpoint: 'hostname1:port1,hostname2:port2,...'
+
+ Then put python value to vineyard.
+
+ .. code:: python
+
+ >>> os.environ['VINEYARD_IPC_SOCKET'] = '/var/run/vineyard.sock'
+ >>> arr = np.arange(8)
+ >>> arr_id = vineyard.put(arr)
+ >>> arr_id
+ 00002ec13bc81226
+
+ Parameters:
+ value:
+ The python value that will be put to vineyard. Supported python value
+ types are decided by modules that registered to vineyard. By default,
+ python value can be put to vineyard after serialized as a bytes buffer
+ using pickle.
+ builder: optional
+ When putting python value to vineyard, an optional *builder* can be
+ specified to tell vineyard how to construct the corresponding vineyard
+ :class:`Object`. If not specified, the default builder context will be
+ used to select a proper builder.
+ persist: bool, optional
+ If true, persist the object after creation.
+ name: str, optional
+ If given, the name will be automatically associated with the resulted
+ object. Note that only take effect when the object is persisted.
+ kw:
+ User-specific argument that will be passed to the builder.
+
+ Returns:
+ ObjectID: The result object id will be returned.
+ """
+
+ client = connect()
+ return client.put(value, builder, persist, name, **kwargs)
+
+
+def get(
+ object_id: Optional[ObjectID] = None,
+ name: Optional[str] = None,
+ resolver: Optional[ResolverContext] = None,
+ fetch: bool = False,
+ **kwargs,
+):
+ """
+ Connect the vineyard server by the following Environment Variables:
+
+ VINEYARD_IPC_SOCKET:
+ UNIX domain socket path to setup an IPC connection.
+ E.g. /var/run/vineyard.sock
+ VINEYARD_RPC_ENDPOINT:
+ TCP endpoint to setup an RPC connection.
+ E.g. 127.0.0.1:9600
+ VINEYARD_CONFIG:
+ Either be a path to a YAML configuration file or a path to a
+ directory containing the default config file `vineyard-config.yaml`.
+
+ The configuration file should be like:
+
+ .. code:: yaml
+
+ Vineyard:
+ IPCSocket: '/path/to/vineyard.sock'
+ RPCEndpoint: 'hostname1:port1,hostname2:port2,...'
+
+ Then get vineyard object as python value.
+
+ .. code:: python
+
+ >>> os.environ['VINEYARD_IPC_SOCKET'] = '/var/run/vineyard.sock'
+ >>> arr_id = vineyard.ObjectID('00002ec13bc81226')
+ >>> arr = vineyard.get(arr_id)
+ >>> arr
+ array([0, 1, 2, 3, 4, 5, 6, 7])
+
+ Parameters:
+ object_id: ObjectID
+ The object id that will be obtained from vineyard.
+ name: ObjectID
+ The object name that will be obtained from vineyard, ignored if
+ ``object_id`` is not None.
+ resolver:
+ When retrieving vineyard object, an optional *resolver* can be specified.
+ If no resolver given, the default resolver context will be used.
+ fetch:
+ Whether to trigger a migration when the target object is located on
+ remote instances.
+ kw:
+ User-specific argument that will be passed to the builder.
+
+ Returns:
+ A python object that return by the resolver, by resolving an vineyard object.
+ """
+
+ client = connect()
+ return client.get(object_id, name, resolver, fetch, **kwargs)
+
+#! /usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2020-2023 Alibaba Group Holding Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import contextlib
+import copy
+import inspect
+import threading
+import warnings
+from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import Generator
+from typing import Optional
+
+from vineyard._C import IPCClient
+from vineyard._C import Object
+from vineyard._C import ObjectID
+from vineyard._C import ObjectMeta
+from vineyard._C import RPCClient
+
+
+[docs]class BuilderContext:
+ def __init__(self, parent_context: Optional["BuilderContext"] = None):
+ self._factory = dict()
+ if parent_context is not None:
+ self._parent_context = parent_context
+ else:
+ self._parent_context = self
+
+ def __str__(self) -> str:
+ return str(self._factory)
+
+ @property
+ def parent_context(self) -> "BuilderContext":
+ return self._parent_context
+
+[docs] def register(self, type_id: type, builder: Callable):
+ """Register a Python type to the builder context.
+
+ Parameters
+ ----------
+ type_id: Python type
+ Like `int`, or `numpy.ndarray`
+
+ builder: callable, e.g., a method, callable object
+ A builder translates a python object to vineyard, it accepts a Python
+ value as parameter, and returns an vineyard object as result.
+ """
+ self._factory[type_id] = builder
+
+[docs] def run(self, client, value, **kw):
+ """Follows the MRO to find the proper builder for given python value.
+
+ Here "Follows the MRO" implies:
+
+ - If the type of python value has been found in the context, the registered
+ builder will be used.
+
+ - If not, it follows the MRO chain from down to top to find a registered
+ Python type and used the associated builder.
+
+ - When the traversal reaches the :code:`object` type, since there's a default
+ builder that serialization the python value, the parameter will be serialized
+ and be put into a blob.
+ """
+
+ # if the python value comes from a vineyard object, we choose to just reuse it.
+
+ # N.B.: don't do that.
+ #
+ # we still construct meta, and optimize the duplicate copy in blob builder
+ #
+ # base = getattr(value, '__vineyard_ref', None)
+ # if base:
+ # return base.meta
+
+ for ty in type(value).__mro__:
+ if ty in self._factory:
+ builder_func_sig = inspect.getfullargspec(self._factory[ty])
+ if (
+ 'builder' in builder_func_sig.args
+ or builder_func_sig.varkw is not None
+ ):
+ kw['builder'] = self
+ return self._factory[ty](client, value, **kw)
+ raise RuntimeError('Unknown type to build as vineyard object')
+
+ def __call__(self, client, value, **kw):
+ return self.run(client, value, **kw)
+
+ def extend(self, builders=None):
+ builder = BuilderContext(self)
+ builder._factory = copy.copy( # pylint: disable=unused-private-member
+ self._factory
+ )
+ if builders:
+ builder._factory.update(builders)
+ return builder
+
+
+default_builder_context = BuilderContext()
+
+_builder_context_local = threading.local()
+_builder_context_local.default_builder = default_builder_context
+
+
+[docs]def get_current_builders() -> BuilderContext:
+ '''Obtain the current builder context.'''
+ default_builder = getattr(_builder_context_local, 'default_builder', None)
+ if not default_builder:
+ default_builder = default_builder_context.extend()
+ return default_builder
+
+
+[docs]@contextlib.contextmanager
+def builder_context(
+ builders: Optional[Dict[type, Callable]] = None,
+ base: Optional[BuilderContext] = None,
+) -> Generator[BuilderContext, Any, Any]:
+ """Open a new context for register builders, without populating outside global
+ environment.
+
+ See Also:
+ resolver_context
+ driver_context
+ """
+ current_builder = get_current_builders()
+ try:
+ builders = builders or dict()
+ base = base or current_builder
+ local_builder = base.extend(builders)
+ _builder_context_local.default_builder = local_builder
+ yield local_builder
+ finally:
+ _builder_context_local.default_builder = current_builder
+
+
+def put(
+ client,
+ value: Any,
+ builder: Optional[BuilderContext] = None,
+ persist: bool = False,
+ name: Optional[str] = None,
+ **kwargs
+):
+ """Put python value to vineyard.
+
+ .. code:: python
+
+ >>> arr = np.arange(8)
+ >>> arr_id = client.put(arr)
+ >>> arr_id
+ 00002ec13bc81226
+
+ Parameters:
+ client: IPCClient or RPCClient
+ The vineyard client to use.
+ value:
+ The python value that will be put to vineyard. Supported python value
+ types are decided by modules that registered to vineyard. By default,
+ python value can be put to vineyard after serialized as a bytes buffer
+ using pickle.
+ builder: optional
+ When putting python value to vineyard, an optional *builder* can be
+ specified to tell vineyard how to construct the corresponding vineyard
+ :class:`Object`. If not specified, the default builder context will be
+ used to select a proper builder.
+ persist: bool, optional
+ If true, persist the object after creation.
+ name: str, optional
+ If given, the name will be automatically associated with the resulted
+ object. Note that only take effect when the object is persisted.
+ kw:
+ User-specific argument that will be passed to the builder.
+
+ Returns:
+ ObjectID: The result object id will be returned.
+ """
+ if builder is not None:
+ return builder(client, value, **kwargs)
+
+ meta = get_current_builders().run(client, value, **kwargs)
+
+ # the builders is expected to return an :class:`ObjectMeta`, or an
+ # :class:`Object` (in the `bytes_builder` and `memoryview` builder).
+ if isinstance(meta, (ObjectMeta, Object)):
+ r = meta.id
+ else:
+ r = meta # None or ObjectID
+
+ if isinstance(r, ObjectID):
+ if persist:
+ client.persist(r)
+ if name is not None:
+ if persist:
+ client.put_name(r, name)
+ else:
+ warnings.warn(
+ "The object is not persisted, the name won't be associated.",
+ UserWarning,
+ )
+ return r
+
+
+setattr(IPCClient, 'put', put)
+setattr(RPCClient, 'put', put)
+
+__all__ = [
+ 'default_builder_context',
+ 'builder_context',
+ 'get_current_builders',
+]
+
+#! /usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2020-2023 Alibaba Group Holding Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import contextlib
+import os
+import warnings
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import as_completed
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
+from vineyard import envvars
+from vineyard._C import Blob
+from vineyard._C import BlobBuilder
+from vineyard._C import IPCClient
+from vineyard._C import NotEnoughMemoryException
+from vineyard._C import Object
+from vineyard._C import ObjectID
+from vineyard._C import ObjectMeta
+from vineyard._C import RemoteBlob
+from vineyard._C import RemoteBlobBuilder
+from vineyard._C import RPCClient
+from vineyard._C import VineyardException
+from vineyard._C import _connect
+from vineyard.core.builder import BuilderContext
+from vineyard.core.builder import put
+from vineyard.core.resolver import ResolverContext
+from vineyard.core.resolver import get
+
+
+def _apply_docstring(func):
+ def _apply(fn):
+ fn.__doc__ = func.__doc__
+ return fn
+
+ return _apply
+
+
+def _parse_configuration(config) -> Tuple[Optional[str], Optional[str]]:
+ '''Parse vineyard IPC socket and RPC endpoints from configuration.
+
+ Parameters:
+ config: Path to a YAML configuration file or a directory containing
+ the default config file `vineyard-config.yaml`. If you define
+ the directory of vineyard config, the
+ `/directory/vineyard-config.yaml` and
+ `/directory/vineyard/vineyard-config.yaml` will be parsed.
+
+ Returns:
+ (socket, endpoints): IPC socket path and RPC endpoints.
+ '''
+ if not config:
+ return None, None
+
+ try:
+ import yaml # pylint: disable=import-outside-toplevel
+ except ImportError:
+ return None, None
+
+ if os.path.isdir(config):
+ config_options = ['vineyard-config.yaml', 'vineyard/vineyard-config.yaml']
+ for config_option in config_options:
+ config_path = os.path.join(config, config_option)
+ if os.path.isfile(config_path):
+ config = config_path
+ break
+ if not os.path.isfile(config):
+ return None, None
+
+ try:
+ with open(config, 'r', encoding='utf-8') as f:
+ vineyard_config = yaml.safe_load(f).get('Vineyard', {})
+ except: # noqa: E722, pylint: disable=bare-except
+ return None, None
+
+ ipc_socket = vineyard_config.get('IPCSocket', None)
+ rpc_endpoint = vineyard_config.get('RPCEndpoint', None)
+
+ if ipc_socket:
+ if not os.path.isabs(ipc_socket):
+ base_dir = os.path.dirname(config) if os.path.isfile(config) else config
+ ipc_socket = os.path.join(base_dir, ipc_socket)
+
+ if not os.path.exists(ipc_socket):
+ ipc_socket = None
+
+ return ipc_socket, rpc_endpoint
+
+
+def _is_blob(object_id: ObjectID):
+ """_is_blob_
+
+ Args:
+ object_id (ObjectID): ObjectID to check if it is a blob
+
+ Returns:
+ bool: True if the object_id is a blob, False otherwise
+ """
+ return int(object_id) & 0x8000000000000000
+
+
+def _traverse_blobs(meta: ObjectMeta, blobs=None):
+ """_traverse_blobs_
+
+ Recursively traverses ObjectMeta to find and accumulate blob IDs by instance_id.
+
+ Args:
+ meta (ObjectMeta): ObjectMeta to traverse for blobs.
+ blobs (dict, optional): Accumulator for blobs organized by instance_id.
+
+ Returns:
+ dict: A dictionary of blobs organized by instance_id.
+ """
+
+ if blobs is None:
+ blobs = {}
+
+ def add_blob(instance_id, blob_id):
+ if instance_id not in blobs:
+ blobs[instance_id] = []
+ blobs[instance_id].append(blob_id)
+
+ if _is_blob(meta.id):
+ add_blob(meta.instance_id, meta.id)
+ else:
+ for _, v in meta.items():
+ if isinstance(v, ObjectMeta):
+ if _is_blob(v.id):
+ add_blob(v.instance_id, v.id)
+ else:
+ _traverse_blobs(v, blobs)
+
+ return blobs
+
+
+[docs]class Client:
+ """Client is responsible for managing IPC and RPC clients for Vineyard
+ and provides a high-level interface to fetch objects from the Vineyard cluster.
+ """
+
+ def __init__(
+ self,
+ socket: str = None,
+ port: Union[int, str] = None,
+ # move host after port to make sure unnamed (host, port) works
+ host: str = None,
+ endpoint: Tuple[str, Union[str, int]] = None,
+ session: int = None,
+ username: str = None,
+ password: str = None,
+ config: str = None,
+ ):
+ """Connects to the vineyard IPC socket and RPC socket.
+
+ - For the IPC Client, the argument `socket` takes precedence over the
+ environment variable `VINEYARD_IPC_SOCKET`, which in turn takes precedence
+ over the `IPCSocket` field in the config file."
+ - For the RPC Client, the argument `endpoint` takes precedence over the
+ argument `host` and `port`, which in turn takes precedence over the
+ environment variable `VINEYARD_RPC_ENDPOINT`, which further takes precedence
+ over the `RPCEndpoint` field in the config file.
+
+ The `connect()` API can be used in following ways:
+
+ - `connect()` without any arguments, which will try to connect to the vineyard
+ by resolving endpoints from the environment variables.
+ - `connect('/path/to/vineyard.sock')`, which will try to establish an IPC
+ connection.
+ - `connect('hostname:port')`, which will try to establish an RPC connection.
+ - `connect('hostname', port)`, which will try to establish an RPC connection.
+ - `connect(endpoint=('hostname', port))`, which will try to establish an RPC
+ connection.
+ - `connect(config='/path/to/vineyard-config.yaml')`, which will try to
+ resolve the IPC socket and RPC endpoints from the configuration file.
+
+ Parameters:
+ socket: Optional, the path to the IPC socket, or RPC endpoints of format
+ `host:port`.
+ port: Optional, the port of the RPC endpoint.
+ host: Optional, the host of the RPC endpoint.
+ endpoint: Optional, the RPC endpoint of format `host:port`.
+ session: Optional, the session id to connect.
+ username: Optional, the required username of vineyardd when authentication
+ is enabled.
+ password: Optional, the required password of vineyardd when authentication
+ is enabled.
+ config: Optional, can either be a path to a YAML configuration file or
+ a path to a directory containing the default config file
+ `vineyard-config.yaml`. Also, the environment variable
+ `VINEYARD_CONFIG` can be used to specify the
+ path to the configuration file. If not defined, the default
+ config file `/var/run/vineyard-config.yaml` or
+ `/var/run/vineyard/vineyard-config.yaml`
+
+ The content of the configuration file should has the following content:
+
+ .. code:: yaml
+
+ Vineyard:
+ IPCSocket: '/path/to/vineyard.sock'
+ RPCEndpoint: 'hostname1:port1,hostname2:port2,...'
+ """
+ self._ipc_client: IPCClient = None
+ self._rpc_client: RPCClient = None
+
+ kwargs = {}
+ if session is not None:
+ kwargs['session'] = session
+ if username is not None:
+ kwargs['username'] = username
+ if password is not None:
+ kwargs['password'] = password
+
+ if socket is not None and port is not None and host is None:
+ socket, host = None, socket
+
+ hosts, ports = [], []
+ if not socket:
+ socket = os.getenv('VINEYARD_IPC_SOCKET', None)
+ if not endpoint and not (host and port):
+ endpoint = os.getenv('VINEYARD_RPC_ENDPOINT', None)
+ if not config:
+ config = os.getenv('VINEYARD_CONFIG', '/var/run')
+ if endpoint:
+ if not isinstance(endpoint, (tuple, list)):
+ for ep in endpoint.split(','):
+ h, p = [e.strip() for e in ep.split(':')]
+ hosts.append(h)
+ ports.append(p)
+ else:
+ h, p = endpoint
+ hosts = [h]
+ ports = [p]
+
+ if host and port:
+ hosts.append(host)
+ ports.append(port)
+
+ if config and ((not socket) or (not (hosts and ports))):
+ ipc_socket, rpc_endpoint = _parse_configuration(config)
+ if ipc_socket and not socket:
+ socket = ipc_socket
+ if rpc_endpoint and not (hosts and ports):
+ for ep in rpc_endpoint.split(','):
+ h, p = [e.strip() for e in ep.split(':')]
+ hosts.append(h)
+ ports.append(p)
+
+ if socket:
+ self._ipc_client = _connect(socket, **kwargs)
+ for host, port in zip(hosts, ports):
+ try:
+ self._rpc_client = _connect(host, port, **kwargs)
+ break
+ except VineyardException:
+ continue
+
+ self._spread = False
+ self._compression = True
+ if self._ipc_client is None and self._rpc_client is None:
+ raise ConnectionError(
+ "Failed to connect to vineyard via both IPC and RPC connection. "
+ "Arguments, environment variables `VINEYARD_IPC_SOCKET` "
+ "and `VINEYARD_RPC_ENDPOINT`, as well as the configuration file, "
+ "are all unavailable."
+ )
+
+ @property
+ def compression(self) -> bool:
+ '''Whether the compression is enabled for underlying RPC client.'''
+ if self._rpc_client:
+ return self._rpc_client.compression
+ return self._compression
+
+ @compression.setter
+ def compression(self, value: bool = True):
+ if self._rpc_client:
+ self._rpc_client.compression = value
+ self._compression = value
+
+ @property
+ def spread(self) -> bool:
+ '''Whether the spread is enabled for underlying RPC client.'''
+ return self._spread
+
+ @spread.setter
+ def spread(self, value: bool = False):
+ self._spread = value
+
+ @property
+ def ipc_client(self) -> IPCClient:
+ assert self._ipc_client is not None, "IPC client is not available."
+ return self._ipc_client
+
+ @property
+ def rpc_client(self) -> RPCClient:
+ assert self._rpc_client is not None, "RPC client is not available."
+ return self._rpc_client
+
+ def has_ipc_client(self):
+ return self._ipc_client is not None
+
+ def has_rpc_client(self):
+ return self._rpc_client is not None
+
+ def default_client(self) -> Union[IPCClient, RPCClient]:
+ return self._ipc_client if self._ipc_client else self._rpc_client
+
+ # The following functions are wrappers of the corresponding functions in the
+ # ClientBase class.
+
+[docs] @_apply_docstring(IPCClient.create_metadata)
+ def create_metadata(
+ self, metadata: Union[ObjectMeta, List[ObjectMeta]], instance_id: int = None
+ ) -> Union[ObjectMeta, List[ObjectMeta]]:
+ if instance_id is not None:
+ return self.default_client().create_metadata(metadata, instance_id)
+ return self.default_client().create_metadata(metadata)
+
+[docs] @_apply_docstring(IPCClient.delete)
+ def delete(
+ self,
+ object: Union[ObjectID, Object, ObjectMeta, List[ObjectID]],
+ force: bool = False,
+ deep: bool = True,
+ memory_trim: bool = False,
+ ) -> None:
+ return self.default_client().delete(object, force, deep, memory_trim)
+
+ @_apply_docstring(IPCClient.create_stream)
+ def create_stream(self, id: ObjectID) -> None:
+ return self.default_client().create_stream(id)
+
+ @_apply_docstring(IPCClient.open_stream)
+ def open_stream(self, id: ObjectID, mode: str) -> None:
+ return self.default_client().open_stream(id, mode)
+
+ @_apply_docstring(IPCClient.push_chunk)
+ def push_chunk(self, stream_id: ObjectID, chunk: ObjectID) -> None:
+ return self.default_client().push_chunk(stream_id, chunk)
+
+ @_apply_docstring(IPCClient.next_chunk_id)
+ def next_chunk_id(self, stream_id: ObjectID) -> ObjectID:
+ return self.default_client().next_chunk_id(stream_id)
+
+ @_apply_docstring(IPCClient.next_chunk_meta)
+ def next_chunk_meta(self, stream_id: ObjectID) -> ObjectMeta:
+ return self.default_client().next_chunk_meta(stream_id)
+
+ @_apply_docstring(IPCClient.next_chunk)
+ def next_chunk(self, stream_id: ObjectID) -> Object:
+ return self.default_client().next_chunk(stream_id)
+
+ @_apply_docstring(IPCClient.stop_stream)
+ def stop_stream(self, stream_id: ObjectID, failed: bool) -> None:
+ return self.default_client().stop_stream(stream_id, failed)
+
+ @_apply_docstring(IPCClient.drop_stream)
+ def drop_stream(self, stream_id: ObjectID) -> None:
+ return self.default_client().drop_stream(stream_id)
+
+[docs] @_apply_docstring(IPCClient.persist)
+ def persist(self, object: Union[ObjectID, Object, ObjectMeta]) -> None:
+ return self.default_client().persist(object)
+
+[docs] @_apply_docstring(IPCClient.exists)
+ def exists(self, object: ObjectID) -> bool:
+ return self.default_client().exists(object)
+
+[docs] @_apply_docstring(IPCClient.shallow_copy)
+ def shallow_copy(
+ self, object_id: ObjectID, extra_metadata: dict = None
+ ) -> ObjectID:
+ if extra_metadata:
+ return self.default_client().shallow_copy(object_id, extra_metadata)
+ return self.default_client().shallow_copy(object_id)
+
+[docs] @_apply_docstring(IPCClient.list_names)
+ def list_names(
+ self, pattern: str, regex: bool = False, limit: int = 5
+ ) -> List[str]:
+ return self.default_client().list_names(pattern, regex, limit)
+
+[docs] @_apply_docstring(IPCClient.put_name)
+ def put_name(self, object: Union[Object, ObjectMeta, ObjectID], name: str) -> None:
+ return self.default_client().put_name(object, name)
+
+[docs] @_apply_docstring(IPCClient.get_name)
+ def get_name(self, name: str, wait: bool = False) -> ObjectID:
+ return self.default_client().get_name(name, wait)
+
+[docs] @_apply_docstring(IPCClient.drop_name)
+ def drop_name(self, name: str) -> None:
+ return self.default_client().drop_name(name)
+
+[docs] @_apply_docstring(IPCClient.sync_meta)
+ def sync_meta(self) -> None:
+ return self.default_client().sync_meta()
+
+ @_apply_docstring(IPCClient.migrate)
+ def migrate(self, object_id: ObjectID) -> ObjectID:
+ return self.default_client().migrate(object_id)
+
+[docs] @_apply_docstring(IPCClient.clear)
+ def clear(self) -> None:
+ return self.default_client().clear()
+
+[docs] @_apply_docstring(IPCClient.memory_trim)
+ def memory_trim(self) -> bool:
+ return self.default_client().memory_trim()
+
+ @_apply_docstring(IPCClient.label)
+ def label(
+ self,
+ object_id: ObjectID,
+ key_or_labels: Union[str, Dict[str, str]],
+ value: str = None,
+ ) -> None:
+ if isinstance(key_or_labels, dict) and value is None:
+ return self.default_client().label(object_id, key_or_labels)
+ else:
+ return self.default_client().label(object_id, key_or_labels, value)
+
+ @_apply_docstring(IPCClient.evict)
+ def evict(self, objects: List[ObjectID]) -> None:
+ return self.default_client().evict(objects)
+
+ @_apply_docstring(IPCClient.load)
+ def load(self, objects: List[ObjectID], pin: bool = False) -> None:
+ return self.default_client().load(objects, pin)
+
+ @_apply_docstring(IPCClient.unpin)
+ def unpin(self, objects: List[ObjectID]) -> None:
+ return self.default_client().unpin(objects)
+
+[docs] @_apply_docstring(IPCClient.reset)
+ def reset(self) -> None:
+ if self._ipc_client:
+ self._ipc_client.reset()
+ if self._rpc_client:
+ self._rpc_client.reset()
+
+ @property
+ @_apply_docstring(IPCClient.connected)
+ def connected(self):
+ return self.default_client().connected
+
+ @property
+ @_apply_docstring(IPCClient.instance_id)
+ def instance_id(self):
+ return self.default_client().instance_id
+
+ @property
+ @_apply_docstring(IPCClient.meta)
+ def meta(self):
+ return self.default_client().meta
+
+ @property
+ @_apply_docstring(IPCClient.status)
+ def status(self):
+ return self.default_client().status
+
+ @_apply_docstring(IPCClient.debug)
+ def debug(self, debug: dict):
+ return self.default_client().debug(debug)
+
+ @property
+ @_apply_docstring(IPCClient.ipc_socket)
+ def ipc_socket(self):
+ return self.default_client().ipc_socket
+
+ @property
+ @_apply_docstring(IPCClient.rpc_endpoint)
+ def rpc_endpoint(self):
+ if self._rpc_client:
+ return self._rpc_client.rpc_endpoint
+ return self.default_client().rpc_endpoint
+
+ @property
+ @_apply_docstring(IPCClient.is_ipc)
+ def is_ipc(self):
+ return self.default_client().is_ipc
+
+ @property
+ @_apply_docstring(IPCClient.is_rpc)
+ def is_rpc(self):
+ return self.default_client().is_rpc
+
+ @property
+ @_apply_docstring(IPCClient.version)
+ def version(self):
+ return self.default_client().version
+
+ # The following functions are wrappers of the corresponding functions in the
+ # IPCClient and RPCClient classes.
+
+[docs] @_apply_docstring(IPCClient.create_blob)
+ def create_blob(
+ self, size: Union[int, List[int]]
+ ) -> Union[BlobBuilder, List[BlobBuilder]]:
+ return self.ipc_client.create_blob(size)
+
+[docs] @_apply_docstring(IPCClient.create_empty_blob)
+ def create_empty_blob(self) -> BlobBuilder:
+ return self.ipc_client.create_empty_blob()
+
+[docs] @_apply_docstring(IPCClient.get_blob)
+ def get_blob(self, object_id: ObjectID, unsafe: bool = False) -> Blob:
+ return self.ipc_client.get_blob(object_id, unsafe)
+
+[docs] @_apply_docstring(IPCClient.get_blobs)
+ def get_blobs(self, object_ids: List[ObjectID], unsafe: bool = False) -> List[Blob]:
+ return self.ipc_client.get_blobs(object_ids, unsafe)
+
+[docs] @_apply_docstring(RPCClient.create_remote_blob)
+ def create_remote_blob(
+ self, blob_builder: Union[RemoteBlobBuilder, List[RemoteBlobBuilder]]
+ ) -> Union[ObjectMeta, List[ObjectMeta]]:
+ return self.rpc_client.create_remote_blob(blob_builder)
+
+[docs] @_apply_docstring(RPCClient.get_remote_blob)
+ def get_remote_blob(self, object_id: ObjectID, unsafe: bool = False) -> RemoteBlob:
+ return self.rpc_client.get_remote_blob(object_id, unsafe)
+
+[docs] @_apply_docstring(RPCClient.get_remote_blobs)
+ def get_remote_blobs(
+ self, object_ids: List[ObjectID], unsafe: bool = False
+ ) -> List[RemoteBlob]:
+ return self.rpc_client.get_remote_blobs(object_ids, unsafe)
+
+[docs] @_apply_docstring(IPCClient.get_object)
+ def get_object(
+ self, object_id: ObjectID, sync_remote: bool = True, fetch: bool = False
+ ) -> Object:
+ """
+ Fetches the object associated with the given object_id from Vineyard.
+ The IPC client is preferred if it's available, otherwise the RPC client
+ """
+ return self._fetch_object(
+ object_id, sync_remote=sync_remote, enable_migrate=fetch
+ )
+
+[docs] @_apply_docstring(IPCClient.get_objects)
+ def get_objects(
+ self, object_ids: List[ObjectID], sync_remote: bool = True
+ ) -> List[Object]:
+ objects = []
+ for object_id in object_ids:
+ objects.append(self.get_object(object_id, sync_remote))
+ return objects
+
+[docs] @_apply_docstring(IPCClient.get_meta)
+ def get_meta(
+ self,
+ object_id: ObjectID,
+ sync_remote: bool = True,
+ ) -> ObjectMeta:
+ return self.default_client().get_meta(object_id, sync_remote)
+
+[docs] @_apply_docstring(IPCClient.get_metas)
+ def get_metas(
+ self, object_ids: List[ObjectID], sync_remote: bool = True
+ ) -> List[ObjectMeta]:
+ metas = []
+ for object_id in object_ids:
+ metas.append(self.get_meta(object_id, sync_remote))
+ return metas
+
+[docs] @_apply_docstring(IPCClient.list_objects)
+ def list_objects(
+ self, pattern: str, regex: bool = False, limit: int = 5
+ ) -> List[ObjectID]:
+ return self.default_client().list_objects(pattern, regex, limit)
+
+[docs] @_apply_docstring(IPCClient.list_metadatas)
+ def list_metadatas(
+ self, pattern: str, regex: bool = False, limit: int = 5, nobuffer: bool = False
+ ) -> List[ObjectMeta]:
+ return self.default_client().list_metadatas(pattern, regex, limit, nobuffer)
+
+ @_apply_docstring(IPCClient.new_buffer_chunk)
+ def new_buffer_chunk(self, stream: ObjectID, size: int) -> memoryview:
+ return self.ipc_client.new_buffer_chunk(stream, size)
+
+ @_apply_docstring(IPCClient.next_buffer_chunk)
+ def next_buffer_chunk(self, stream: ObjectID) -> memoryview:
+ return self.ipc_client.next_buffer_chunk(stream)
+
+[docs] @_apply_docstring(IPCClient.allocated_size)
+ def allocated_size(self, object_id: Union[Object, ObjectID]) -> int:
+ return self.ipc_client.allocated_size(object_id)
+
+
+
+
+
+ @property
+ @_apply_docstring(RPCClient.remote_instance_id)
+ def remote_instance_id(self) -> int:
+ return self.rpc_client.remote_instance_id
+
+[docs] @_apply_docstring(IPCClient.close)
+ def close(self) -> None:
+ if self._ipc_client:
+ self._ipc_client.close()
+ if self._rpc_client:
+ self._rpc_client.close()
+
+ @_apply_docstring(IPCClient.fork)
+ def fork(self) -> 'Client':
+ if self._ipc_client:
+ self._ipc_client = self._ipc_client.fork()
+ if self._rpc_client:
+ self._rpc_client = self._rpc_client.fork()
+ return self
+
+ def _fetch_object(
+ self, object_id: ObjectID, enable_migrate: bool, sync_remote: bool = True
+ ) -> Object:
+ meta = self.get_meta(object_id, sync_remote=sync_remote)
+
+ if self.has_ipc_client() and enable_migrate:
+ # no need to sync remote metadata as the metadata is already fetched
+ return self._ipc_client.get_object(object_id, fetch=True, sync_remote=False)
+
+ blobs = _traverse_blobs(meta)
+
+ # If the object is local, return it directly
+ if blobs.keys() == {self.instance_id}:
+ return Object.from_(meta)
+
+ cluster_info = self.default_client().meta
+ meta.force_local()
+ meta._client = None
+
+ with ThreadPoolExecutor() as executor:
+ futures = {
+ executor.submit(
+ self._fetch_blobs_from_instance,
+ cluster_info,
+ instance_id,
+ blobs[instance_id],
+ self.compression,
+ )
+ for instance_id in blobs
+ if instance_id != self.instance_id
+ }
+
+ for future in as_completed(futures):
+ fetched_blobs = future.result()
+ for blob in fetched_blobs:
+ meta.add_remote_blob(blob)
+
+ return Object.from_(meta)
+
+ def _fetch_blobs_from_instance(
+ self, cluster_info, instance_id, blob_ids, compression
+ ) -> Object:
+ """Fetches all blobs from a given instance id in the Vineyard cluster.
+
+ Args:
+ cluster_info (Dict): The cluster information of the Vineyard cluster.
+ instance_id (int): The instance id to fetch blobs from.
+ blob_ids (List): The list of blob ids to fetch.
+ compression (bool): Whether to enable compression for RPC Client.
+
+ Returns:
+ RemoteBlob(List): The list of fetched remote blobs.
+ """
+ instance_status = cluster_info.get(instance_id)
+
+ if instance_status is None or instance_status['rpc_endpoint'] is None:
+ raise RuntimeError(
+ "The rpc endpoint of the vineyard instance "
+ f"{instance_id} is not available."
+ )
+
+ if self.has_rpc_client() and self.remote_instance_id == instance_id:
+ remote_client = self._rpc_client
+ else:
+ host, port = instance_status['rpc_endpoint'].split(':')
+ try:
+ with envvars('VINEYARD_RPC_SKIP_RETRY', '1'):
+ remote_client = _connect(host, port)
+ remote_client.compression = compression
+ except Exception as exec:
+ raise RuntimeError(
+ f"Failed to connect to the vineyard instance {instance_id} "
+ f"at {host}:{port}."
+ ) from exec
+
+ return remote_client.get_remote_blobs(blob_ids)
+
+ def _connect_and_get_memory(self, instance_id, rpc_endpoint):
+ host, port = rpc_endpoint.split(':')
+ try:
+ new_client = _connect(host, port)
+ current_available_memory = (
+ new_client.status.memory_limit - new_client.status.memory_usage
+ )
+ return instance_id, current_available_memory
+ except Exception:
+ return instance_id, float('-inf')
+
+ def _find_the_most_available_memory_instance(self) -> int:
+ """
+ Find the vineyard instance with the most available memory.
+
+ Returns:
+ int: The instance id with the most available memory.
+ if only have one instance, return -1
+ """
+ futures = []
+ cluster_info = self.default_client().meta
+ with ThreadPoolExecutor() as executor:
+ for instance_id, status in cluster_info.items():
+ if not (
+ status['ipc_socket'] == self.ipc_socket
+ and status['rpc_endpoint'] == self.rpc_endpoint
+ ):
+ futures.append(
+ executor.submit(
+ self._connect_and_get_memory,
+ instance_id,
+ status['rpc_endpoint'],
+ )
+ )
+
+ instance_id_with_most_available_memory = -1
+ available_memory = float('-inf')
+
+ for future in as_completed(futures):
+ instance_id, current_available_memory = future.result()
+ if current_available_memory > available_memory:
+ instance_id_with_most_available_memory = instance_id
+ available_memory = current_available_memory
+
+ return instance_id_with_most_available_memory
+
+[docs] @_apply_docstring(get)
+ def get(
+ self,
+ object_id: Optional[ObjectID] = None,
+ name: Optional[str] = None,
+ resolver: Optional[ResolverContext] = None,
+ fetch: bool = False,
+ **kwargs,
+ ):
+ return get(self, object_id, name, resolver, fetch, **kwargs)
+
+[docs] @_apply_docstring(put)
+ def put(
+ self,
+ value: Any,
+ builder: Optional[BuilderContext] = None,
+ persist: bool = False,
+ name: Optional[str] = None,
+ **kwargs,
+ ):
+ try:
+ return put(self, value, builder, persist, name, **kwargs)
+ except NotEnoughMemoryException as exec:
+ with envvars(
+ {'VINEYARD_RPC_SKIP_RETRY': '1', 'VINEYARD_IPC_SKIP_RETRY': '1'}
+ ):
+ instance_id = self._find_the_most_available_memory_instance()
+ previous_compression_state = self.compression
+ if instance_id == -1:
+ warnings.warn("No other vineyard instance available")
+ raise exec
+ else:
+ meta = self.default_client().meta
+ warnings.warn(
+ f"Put object to the vineyard instance {instance_id}"
+ "with the most available memory."
+ )
+ # connect to the instance with the most available memory
+ self._ipc_client = None
+ if os.path.exists(meta[instance_id]['ipc_socket']):
+ self._ipc_client = _connect(meta[instance_id]['ipc_socket'])
+ # avoid the case the vineyard instance is restarted
+ if self._ipc_client.instance_id != instance_id:
+ self._ipc_client = None
+ host, port = meta[instance_id]['rpc_endpoint'].split(':')
+ self._rpc_client = _connect(host, port)
+ self.compression = previous_compression_state
+ return put(self, value, builder, persist, name, **kwargs)
+
+[docs] @contextlib.contextmanager
+ def with_compression(self, enabled: bool = True):
+ """Disable compression for the following put operations."""
+ compression = self.compression
+ self.compression = enabled
+ yield
+ self.compression = compression
+
+[docs] @contextlib.contextmanager
+ def with_spread(self, enabled: bool = True):
+ """Enable spread for the following put operations."""
+ tmp_spread = self._spread
+ self.spread = enabled
+ yield
+ self.spread = tmp_spread
+
+
+__all__ = ['Client']
+
+#! /usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2020-2023 Alibaba Group Holding Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import contextlib
+import copy
+import functools
+import threading
+
+from sortedcontainers import SortedDict
+
+from vineyard.core.utils import find_most_precise_match
+
+
+[docs]class DriverContext:
+ def __init__(self):
+ self._factory = SortedDict(dict)
+
+ def __str__(self) -> str:
+ return str(self._factory)
+
+ def register(self, typename_prefix, meth, func):
+ if typename_prefix not in self._factory:
+ self._factory[typename_prefix] = dict()
+ self._factory[typename_prefix][meth] = func
+
+ def resolve(self, obj, typename):
+ prefix, methods = find_most_precise_match(typename, self._factory)
+ if prefix:
+ for meth, func in methods.items():
+ # if shouldn't failed, since it has already been wrapped in during
+ # resolving
+ setattr(obj, meth, functools.partial(func, obj))
+ return obj
+
+ def __call__(self, obj, typename):
+ return self.resolve(obj, typename)
+
+ def extend(self, drivers=None):
+ driver = DriverContext()
+ driver._factory.update(((k, copy.copy(v)) for k, v in self._factory.items()))
+ if drivers:
+ for ty, methods in drivers.items():
+ if ty not in self._factory:
+ driver._factory[ty] = dict()
+ driver._factory[ty].update(methods)
+ return driver
+
+
+default_driver_context = DriverContext()
+
+_driver_context_local = threading.local()
+_driver_context_local.default_driver = default_driver_context
+
+
+[docs]def get_current_drivers():
+ '''Obtain current driver context.'''
+ default_driver = getattr(_driver_context_local, 'default_driver', None)
+ if not default_driver:
+ default_driver = default_driver_context.extend()
+ return default_driver
+
+
+[docs]@contextlib.contextmanager
+def driver_context(drivers=None, base=None):
+ """Open a new context for register drivers, without populting outside global
+ environment.
+
+ See Also:
+ builder_context
+ resolver_context
+ """
+ current_driver = get_current_drivers()
+ try:
+ drivers = drivers or dict()
+ base = base or current_driver
+ local_driver = base.extend(drivers)
+ _driver_context_local.default_driver = local_driver
+ yield local_driver
+ finally:
+ _driver_context_local.default_driver = current_driver
+
+
+def register_builtin_drivers(ctx):
+ assert isinstance(ctx, DriverContext)
+
+ # TODO
+ # there's no builtin drivers yet.
+
+
+def registerize(func):
+ """Registerize a method, add a `_factory` attribute and a `register`
+ interface to a method.
+
+ multiple-level register is automatically supported, users can
+
+ >>> open.register(local_io_adaptor)
+ >>> open.register(oss_io_adaptor)
+
+ OR
+
+ >>> open.register('file', local_io_adaptor)
+ >>> open.register('odps', odps_io_adaptor)
+
+ OR
+
+ >>> open.register('file', 'csv', local_csv_reader)
+ >>> open.register('file', 'tsv', local_tsv_reader)
+ """
+
+ @functools.wraps(func)
+ def wrap(*args, **kwargs):
+ return func(*args, **kwargs)
+
+ setattr(wrap, '_factory', None)
+
+ def register(*args):
+ if len(args) == 1:
+ if wrap._factory is None:
+ wrap._factory = []
+ if not isinstance(wrap._factory, list):
+ raise RuntimeError(
+ 'Invalid arguments: inconsistent with existing registerations'
+ )
+ wrap._factory.append(args[0])
+ else:
+ if wrap._factory is None:
+ wrap._factory = {}
+ if not isinstance(wrap._factory, dict):
+ raise RuntimeError(
+ 'Invalid arguments: inconsistent with existing registerations'
+ )
+ root = wrap._factory
+ for arg in args[:-2]:
+ if arg not in root:
+ root[arg] = dict()
+ root = root[arg]
+ if args[-2] not in root:
+ root[args[-2]] = list()
+ root[args[-2]].append(args[-1])
+
+ setattr(wrap, 'register', register)
+
+ return wrap
+
+
+__all__ = [
+ 'default_driver_context',
+ 'register_builtin_drivers',
+ 'driver_context',
+ 'get_current_drivers',
+ 'registerize',
+]
+
+#! /usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2020-2023 Alibaba Group Holding Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import contextlib
+import inspect
+import threading
+from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import Generator
+from typing import Optional
+
+from sortedcontainers import SortedDict
+
+from vineyard._C import IPCClient
+from vineyard._C import Object
+from vineyard._C import ObjectID
+from vineyard._C import RPCClient
+from vineyard.core.utils import find_most_precise_match
+
+
+[docs]class ResolverContext:
+ def __init__(self, parent_context: Optional["ResolverContext"] = None):
+ self._factory = SortedDict()
+ if parent_context is not None:
+ self._parent_context = parent_context
+ else:
+ self._parent_context = self
+
+ def __str__(self) -> str:
+ return str(self._factory)
+
+ def __repr__(self) -> str:
+ return repr(self._factory)
+
+ @property
+ def parent_context(self) -> "ResolverContext":
+ return self._parent_context
+
+ def register(self, typename_prefix: str, resolver: Callable):
+ self._factory[typename_prefix] = resolver
+
+ def run(self, obj: Any, **kw):
+ typename = obj.meta.typename
+ prefix, resolver = find_most_precise_match(typename, self._factory)
+ vineyard_client = kw.pop('__vineyard_client', None)
+ if prefix:
+ resolver_func_sig = inspect.getfullargspec(resolver)
+ if resolver_func_sig.varkw is not None:
+ value = resolver(obj, resolver=self, **kw)
+ else:
+ try:
+ # don't pass the `**kw`.
+ if 'resolver' in resolver_func_sig.args:
+ value = resolver(obj, resolver=self)
+ else:
+ value = resolver(obj)
+ except Exception as e:
+ raise RuntimeError( # pylint: disable=raise-missing-from
+ 'Unable to construct the object using resolver: '
+ 'typename is %s, resolver is %s' % (obj.meta.typename, resolver)
+ ) from e
+ if value is None:
+ # if the obj has been resolved by pybind types, and there's no proper
+ # resolver, it shouldn't be an error
+ if type(obj) is not Object: # pylint: disable=unidiomatic-typecheck
+ return obj
+
+ # we might `client.put(None)`
+ return None
+
+ # associate a reference to the base C++ object
+ try:
+ setattr(value, '__vineyard_ref', obj)
+ setattr(value, '__vineyard_client', vineyard_client)
+
+ # register methods
+ from vineyard.core.driver import get_current_drivers
+
+ get_current_drivers().resolve(value, obj.typename)
+ except AttributeError:
+ pass
+
+ return value
+ # keep it as it is
+ return obj
+
+ def __call__(self, obj, **kw):
+ return self.run(obj, **kw)
+
+ def extend(self, resolvers=None):
+ resolver = ResolverContext(self)
+ resolver._factory = ( # pylint: disable=unused-private-member
+ self._factory.copy()
+ )
+ if resolvers:
+ resolver._factory.update(resolvers)
+ return resolver
+
+
+default_resolver_context = ResolverContext()
+
+_resolver_context_local = threading.local()
+_resolver_context_local.default_resolver = default_resolver_context
+
+
+[docs]def get_current_resolvers() -> ResolverContext:
+ '''Obtain current resolver context.'''
+ default_resolver = getattr(_resolver_context_local, 'default_resolver', None)
+ if not default_resolver:
+ default_resolver = default_resolver_context.extend()
+ return default_resolver
+
+
+[docs]@contextlib.contextmanager
+def resolver_context(
+ resolvers: Optional[Dict[str, Callable]] = None,
+ base: Optional[ResolverContext] = None,
+) -> Generator[ResolverContext, Any, Any]:
+ """Open a new context for register resolvers, without populating outside
+ the global environment.
+
+ The :code:`resolver_context` can be useful when users have more than
+ more resolver for a certain type, e.g., the :code:`vineyard::Tensor`
+ object can be resolved as :code:`numpy.ndarray` or :code:`xgboost::DMatrix`.
+
+ We could have
+
+ .. code:: python
+
+ def numpy_resolver(obj):
+ ...
+
+ default_resolver_context.register('vineyard::Tensor', numpy_resolver)
+
+ and
+
+ .. code:: python
+
+ def xgboost_resolver(obj):
+ ...
+
+ default_resolver_context.register('vineyard::Tensor', xgboost_resolver)
+
+ Obviously there's a conflict, and the stackable :code:`resolver_context` could
+ help there,
+
+ .. code:: python
+
+ with resolver_context({'vineyard::Tensor', xgboost_resolver}):
+ ...
+
+ Assuming the default context resolves :code:`vineyard::Tensor` to
+ :code:`numpy.ndarray`, inside the :code:`with resolver_context` the
+ :code:`vineyard::Tensor` will be resolved to :code:`xgboost::DMatrix`,
+ and after exiting the context the global environment will be restored
+ back as default.
+
+ The :code:`with resolver_context` is nestable as well.
+
+ See Also:
+ builder_context
+ driver_context
+ """
+ current_resolver = get_current_resolvers()
+ try:
+ resolvers = resolvers or dict()
+ base = base or current_resolver
+ local_resolver = base.extend(resolvers)
+ _resolver_context_local.default_resolver = local_resolver
+ yield local_resolver
+ finally:
+ _resolver_context_local.default_resolver = current_resolver
+
+
+def get(
+ client,
+ object_id: Optional[ObjectID] = None,
+ name: Optional[str] = None,
+ resolver: Optional[ResolverContext] = None,
+ fetch: bool = False,
+ **kwargs
+):
+ """Get vineyard object as python value.
+
+ .. code:: python
+
+ >>> arr_id = vineyard.ObjectID('00002ec13bc81226')
+ >>> arr = client.get(arr_id)
+ >>> arr
+ array([0, 1, 2, 3, 4, 5, 6, 7])
+
+ Parameters:
+ client: IPCClient or RPCClient
+ The vineyard client to use.
+ object_id: ObjectID
+ The object id that will be obtained from vineyard.
+ name: ObjectID
+ The object name that will be obtained from vineyard, ignored if
+ ``object_id`` is not None.
+ resolver:
+ When retrieving vineyard object, an optional *resolver* can be specified.
+ If no resolver given, the default resolver context will be used.
+ fetch:
+ Whether to trigger a migration when the target object is located on
+ remote instances.
+ kw:
+ User-specific argument that will be passed to the builder.
+
+ Returns:
+ A python object that return by the resolver, by resolving an vineyard object.
+ """
+ # wrap object_id
+ if object_id is not None:
+ if isinstance(object_id, (int, str)):
+ object_id = ObjectID(object_id)
+ elif name is not None:
+ object_id = client.get_name(name)
+
+ obj = client.get_object(object_id, fetch=fetch)
+
+ if resolver is None:
+ resolver = get_current_resolvers()
+ return resolver(obj, __vineyard_client=client, **kwargs)
+
+
+setattr(IPCClient, 'get', get)
+setattr(RPCClient, 'get', get)
+
+__all__ = [
+ 'default_resolver_context',
+ 'resolver_context',
+ 'get_current_resolvers',
+]
+
+#! /usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2020-2023 Alibaba Group Holding Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+from typing import Any
+
+import vineyard
+
+
+def _find_vineyard_socket(path):
+ current_dir = path
+ while current_dir != "/":
+ socket_path = os.path.join(current_dir, "vineyard.sock")
+ if os.path.exists(socket_path):
+ return socket_path
+ current_dir = os.path.dirname(current_dir)
+ return None
+
+
+[docs]def write(
+ value: Any,
+ path: str,
+):
+ """
+ Write python value to vineyard.
+ Notice, the API is only used for CSI driver.
+
+ Parameters:
+ path: str
+ The path that represents a vineyard object.
+
+ .. code:: python
+
+ >>> arr = np.arange(8)
+ >>> vineyard.write(arr)
+ """
+ socket_path = _find_vineyard_socket(path)
+
+ if socket_path is None:
+ raise FileNotFoundError(
+ f"The given path is not generated by vineyard CSI driver: {path}"
+ )
+
+ client = vineyard.connect(socket_path)
+ client.put(value, persist=True, name=path)
+
+
+[docs]def read(
+ path: str,
+):
+ """
+ Read vineyard object from path, and return python value.
+ Notice, the API is only used for CSI driver.
+
+ Parameters:
+ path: str
+ The path that represents a vineyard object.
+
+ .. code:: python
+
+ >>> arr = vineyard.read('/a/b/c/d/f')
+ >>> arr
+ array([0, 1, 2, 3, 4, 5, 6, 7])
+
+ Returns:
+ A python object that return by the resolver, by resolving an vineyard object.
+ """
+ socket_path = _find_vineyard_socket(path)
+
+ if socket_path is None:
+ raise FileNotFoundError(
+ f"The given path is not generated by vineyard CSI driver: {path}"
+ )
+
+ client = vineyard.connect(socket_path)
+ return client.get(name=path)
+
+#! /usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2020-2023 Alibaba Group Holding Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import contextlib
+import logging
+import subprocess
+import sys
+import textwrap
+import time
+
+import pkg_resources
+
+from vineyard.deploy.utils import ssh_base_cmd
+from vineyard.deploy.utils import start_etcd
+
+logger = logging.getLogger('vineyard')
+
+
+[docs]@contextlib.contextmanager
+def start_vineyardd(
+ hosts=None,
+ etcd_endpoints=None,
+ vineyardd_path=None,
+ size='',
+ socket='/var/run/vineyard.sock',
+ rpc_socket_port=9600,
+ debug=False,
+):
+ """Launch a local vineyard cluster in a distributed fashion.
+
+ Parameters:
+ hosts: list of str
+ A list of machines to launch vineyard server.
+ etcd_endpoint: str
+ Launching vineyard using specified etcd endpoints. If not specified,
+ vineyard will launch its own etcd instance.
+ vineyardd_path: str
+ Location of vineyard server program. If not specified, vineyard will
+ use its own bundled vineyardd binary.
+ size: int
+ The memory size limit for vineyard's shared memory. The memory size
+ can be a plain integer or as a fixed-point number using one of these
+ suffixes:
+
+ .. code::
+
+ E, P, T, G, M, K.
+
+ You can also use the power-of-two equivalents: Ei, Pi, Ti, Gi, Mi, Ki.
+ Defaults to "", means not limited.
+
+ For example, the following represent roughly the same value:
+
+ .. code::
+
+ 128974848, 129k, 129M, 123Mi, 1G, 10Gi, ...
+ socket: str
+ The UNIX domain socket socket path that vineyard server will listen on.
+ rpc_socket_port: int
+ The port that vineyard will use to privode RPC service.
+ debug: bool
+ Whether print debug logs.
+ """
+ if vineyardd_path is None:
+ vineyardd_path = pkg_resources.resource_filename('vineyard', 'vineyardd')
+
+ if hosts is None:
+ hosts = ['localhost']
+
+ if etcd_endpoints is None:
+ etcd_ctx = start_etcd(host=hosts[0])
+ _etcd_proc, etcd_endpoints = etcd_ctx.__enter__() # pylint: disable=no-member
+ else:
+ etcd_ctx = None
+
+ env = dict()
+ if debug:
+ env['GLOG_v'] = 11
+
+ command = [
+ vineyardd_path,
+ '--deployment',
+ 'distributed',
+ '--size',
+ str(size),
+ '--socket',
+ socket,
+ '--rpc_socket_port',
+ str(rpc_socket_port),
+ '--etcd_endpoint',
+ etcd_endpoints,
+ ]
+
+ procs = []
+ try:
+ for host in hosts:
+ proc = subprocess.Popen(
+ ssh_base_cmd(host) + command,
+ env=env,
+ stdout=subprocess.PIPE,
+ stderr=sys.__stderr__,
+ universal_newlines=True,
+ encoding='utf-8',
+ )
+ procs.append(proc)
+ time.sleep(1)
+ for proc in procs:
+ rc = proc.poll()
+ if rc is not None:
+ err = textwrap.indent(proc.stdout.read(), ' ' * 4)
+ raise RuntimeError(
+ 'vineyardd exited unexpectedly with '
+ 'code %d: error is:\n%s' % (rc, err)
+ )
+ yield procs, socket, etcd_endpoints
+ finally:
+ logger.info('Distributed vineyardd being killed')
+ for proc in procs:
+ if proc.poll() is None:
+ logger.info('killing the vineyardd')
+ proc.kill()
+ if etcd_ctx is not None:
+ etcd_ctx.__exit__(None, None, None) # pylint: disable=no-member
+
+
+__all__ = ['start_vineyardd']
+
+#! /usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2020-2023 Alibaba Group Holding Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import logging
+import os
+import re
+import tempfile
+import time
+
+try:
+ import kubernetes
+except ImportError:
+ kubernetes = None
+
+from vineyard.deploy.etcd import start_etcd_k8s
+from vineyard.deploy.utils import ensure_kubernetes_namespace
+
+logger = logging.getLogger('vineyard')
+
+
+[docs]def start_vineyardd(
+ namespace='vineyard',
+ size='512Mi',
+ socket='/var/run/vineyard.sock',
+ rpc_socket_port=9600,
+ vineyard_image='vineyardcloudnative/vineyardd:latest',
+ vineyard_image_pull_policy='IfNotPresent',
+ vineyard_image_pull_secrets=None,
+ k8s_client=None,
+):
+ """Launch a vineyard cluster on kubernetes.
+
+ Parameters:
+ namespace: str
+ namespace in kubernetes
+ size: int
+ The memory size limit for vineyard's shared memory. The memory size
+ can be a plain integer or as a fixed-point number using one of these
+ suffixes:
+
+ .. code::
+
+ E, P, T, G, M, K.
+
+ You can also use the power-of-two equivalents: Ei, Pi, Ti, Gi, Mi, Ki.
+
+ For example, the following represent roughly the same value:
+
+ .. code::
+
+ 128974848, 129k, 129M, 123Mi, 1G, 10Gi, ...
+ socket: str
+ The UNIX domain socket socket path that vineyard server will listen on.
+ rpc_socket_port: int
+ The port that vineyard will use to provide RPC service.
+ k8s_client: kubernetes.client.api.ApiClient
+ A kubernetes client. If not specified, vineyard will try to resolve the
+ kubernetes configuration from current context.
+ vineyard_image: str
+ The docker image of vineyardd to launch the daemonset.
+ vineyard_image_pull_policy: str
+ The docker image pull policy of vineyardd.
+ vineyard_image_pull_secrets: str
+ The docker image pull secrets of vineyardd.
+
+ Returns:
+ A list of created kubernetes resources during the deploying process. The
+ resources can be later release using :meth:`delete_kubernetes_objects`.
+
+ See Also:
+ vineyard.deploy.kubernetes.delete_kubernetes_objects
+ """
+ if kubernetes is None:
+ raise RuntimeError('Please install the package python "kubernetes" first')
+
+ if k8s_client is None:
+ kubernetes.config.load_kube_config()
+ k8s_client = kubernetes.client.ApiClient()
+
+ created_objects = []
+ ensure_kubernetes_namespace(namespace, k8s_client=k8s_client)
+ created_objects.extend(start_etcd_k8s(namespace, k8s_client=k8s_client))
+
+ with open(
+ os.path.join(os.path.dirname(__file__), "vineyard.yaml.tpl"),
+ 'r',
+ encoding='utf-8',
+ ) as fp:
+ formatter = {
+ 'Namespace': namespace,
+ 'Size': size,
+ 'Socket': socket,
+ 'Port': rpc_socket_port,
+ 'Image': vineyard_image,
+ 'ImagePullPolicy': vineyard_image_pull_policy,
+ 'ImagePullSecrets': (
+ vineyard_image_pull_secrets if vineyard_image_pull_secrets else 'none',
+ ),
+ }
+ definitions = fp.read().format(**formatter)
+
+ with tempfile.NamedTemporaryFile(
+ mode='w', encoding='utf-8', delete=False
+ ) as rendered:
+ rendered.write(definitions)
+ rendered.flush()
+ rendered.close()
+
+ with open(rendered.name, 'r', encoding='utf-8') as fp:
+ logger.debug(fp.read())
+
+ created_objects.extend(
+ kubernetes.utils.create_from_yaml(
+ k8s_client, rendered.name, namespace=namespace
+ )
+ )
+ return created_objects
+
+
+def recursive_flatten(targets):
+ """Flatten the given maybe nested list as a 1-level list."""
+
+ def _recursive_flatten_impl(destination, targets):
+ if isinstance(targets, (list, tuple)):
+ for target in targets:
+ _recursive_flatten_impl(destination, target)
+ else:
+ destination.append(targets)
+
+ destination = []
+ _recursive_flatten_impl(destination, targets)
+ return destination
+
+
+[docs]def delete_kubernetes_objects(
+ targets, k8s_client=None, verbose=False, wait=False, timeout_seconds=60, **kwargs
+):
+ """Delete the given kubernetes resources.
+
+ Parameters:
+ target: List
+ List of Kubernetes objects
+ k8s_client:
+ The kubernetes client. If not specified, vineyard will try to resolve
+ the kubernetes configuration from current context.
+ verbose: bool
+ Whether to print the deletion logs.
+ wait: bool
+ Whether to wait for the deletion to complete.
+ timeout_seconds: int
+ The timeout in seconds for waiting for the deletion to complete.
+
+ See Also:
+ vineyard.deploy.kubernetes.start_vineyardd
+ vineyard.deploy.kubernetes.delete_kubernetes_object
+ """
+ for target in recursive_flatten(targets):
+ delete_kubernetes_object(
+ target,
+ k8s_client,
+ verbose=verbose,
+ wait=wait,
+ timeout_seconds=timeout_seconds,
+ **kwargs
+ )
+
+
+def delete_kubernetes_object(
+ target, k8s_client=None, verbose=False, wait=False, timeout_seconds=60, **kwargs
+):
+ """Delete the given kubernetes resource.
+
+ Parameters:
+ target: object
+ The Kubernetes objects that will be deleted.
+ k8s_client:
+ The kubernetes client. If not specified, vineyard will try to resolve
+ the kubernetes configuration from current context.
+ verbose: bool
+ If True, print confirmation from the delete action. Defaults to False.
+ wait: bool
+ Whether to wait for the deletion to complete. Defaults to False.
+ timeout_seconds: int
+ The timeout in seconds for waiting for the deletion to complete. Defaults
+ to 60.
+
+ Returns:
+ Status: Return status for calls kubernetes delete method.
+
+ See Also:
+ vineyard.deploy.kubernetes.start_vineyardd
+ vineyard.deploy.kubernetes.delete_kubernetes_objects
+ """
+ if isinstance(target, (list, tuple)):
+ return delete_kubernetes_objects(
+ target, k8s_client, verbose=verbose, wait=wait, **kwargs
+ )
+
+ if kubernetes is None:
+ raise RuntimeError('Please install the package python "kubernetes" first')
+
+ if k8s_client is None:
+ kubernetes.config.load_kube_config()
+ k8s_client = kubernetes.client.ApiClient()
+
+ group, _, version = target.api_version.partition("/")
+ if version == "":
+ version = group
+ group = "core"
+ # Take care for the case e.g. api_type is "apiextensions.k8s.io"
+ # Only replace the last instance
+ group = "".join(group.rsplit(".k8s.io", 1))
+ # convert group name from DNS subdomain format to
+ # python class name convention
+ group = "".join(word.capitalize() for word in group.split("."))
+ fcn_to_call = "{0}{1}Api".format(group, version.capitalize())
+ k8s_api = getattr(kubernetes.client, fcn_to_call)(
+ k8s_client
+ ) # pylint: disable=not-callable
+
+ kind = target.kind
+ kind = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", kind)
+ kind = re.sub("([a-z0-9])([A-Z])", r"\1_\2", kind).lower()
+
+ try:
+ # Expect the user to create namespaced objects more often
+ kwargs["name"] = target.metadata.name
+ if hasattr(k8s_api, "delete_namespaced_{0}".format(kind)):
+ # Decide which namespace we are going to put the object in, if any
+ kwargs["namespace"] = target.metadata.namespace
+ resp = getattr(k8s_api, "delete_namespaced_{0}".format(kind))(**kwargs)
+ else:
+ kwargs.pop("namespace", None)
+ resp = getattr(k8s_api, "delete_{0}".format(kind))(**kwargs)
+ except kubernetes.client.rest.ApiException:
+ # Object already deleted.
+ return None
+ else:
+ # Waiting for delete
+ if wait:
+ start_time = time.time()
+ if hasattr(k8s_api, "read_namespaced_{0}".format(kind)):
+ while True:
+ try:
+ getattr(k8s_api, "read_namespaced_{0}".format(kind))(**kwargs)
+ except kubernetes.client.rest.ApiException as ex:
+ if ex.status != 404:
+ logger.exception(
+ "Deleting %s %s failed", kind, target.metadata.name
+ )
+ break
+ else:
+ time.sleep(1)
+ if time.time() - start_time > timeout_seconds:
+ logger.info(
+ "Deleting %s/%s timeout", kind, target.metadata.name
+ )
+ if verbose:
+ msg = "{0}/{1} deleted.".format(kind, target.metadata.name)
+ if hasattr(resp, "status"):
+ msg += " status='{0}'".format(str(resp.status))
+ logger.info(msg)
+ return resp
+
+
+__all__ = [
+ 'start_vineyardd',
+ 'delete_kubernetes_object',
+ 'delete_kubernetes_objects',
+]
+
+#! /usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2020-2023 Alibaba Group Holding Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import atexit
+import contextlib
+import logging
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+import textwrap
+import time
+from typing import Generator
+from typing import Optional
+from typing import Tuple
+
+from vineyard.deploy.etcd import start_etcd
+from vineyard.deploy.utils import check_socket
+from vineyard.deploy.utils import find_vineyardd_path
+
+logger = logging.getLogger('vineyard')
+
+
+[docs]@contextlib.contextmanager
+def start_vineyardd(
+ meta: Optional[str] = 'etcd',
+ etcd_endpoints: Optional[str] = None,
+ etcd_prefix: Optional[str] = None,
+ vineyardd_path: Optional[str] = None,
+ size: Optional[str] = '',
+ socket: Optional[str] = None,
+ rpc: Optional[str] = True,
+ rpc_socket_port: Optional[str] = 9600,
+ debug=False,
+) -> Generator[Tuple[subprocess.Popen, str, str], None, None]:
+ """Launch a local vineyard cluster.
+
+ Parameters:
+ meta: str, optional.
+ Metadata backend, can be "etcd", "redis" and "local". Defaults to
+ "etcd".
+ etcd_endpoint: str
+ Launching vineyard using specified etcd endpoints. If not specified,
+ vineyard will launch its own etcd instance.
+ etcd_prefix: str
+ Specify a common prefix to establish a local vineyard cluster.
+ vineyardd_path: str
+ Location of vineyard server program. If not specified, vineyard will
+ use its own bundled vineyardd binary.
+ size: int
+ The memory size limit for vineyard's shared memory. The memory size
+ can be a plain integer or as a fixed-point number using one of these
+ suffixes:
+
+ .. code::
+
+ E, P, T, G, M, K.
+
+ You can also use the power-of-two equivalents: Ei, Pi, Ti, Gi, Mi, Ki.
+ Defaults to "", means not limited.
+
+ For example, the following represent roughly the same value:
+
+ .. code::
+
+ 128974848, 129k, 129M, 123Mi, 1G, 10Gi, ...
+ socket: str
+ The UNIX domain socket socket path that vineyard server will listen on.
+ Default is None.
+
+ When the socket parameter is None, a random path under temporary directory
+ will be generated and used.
+ rpc_socket_port: int
+ The port that vineyard will use to provided RPC service.
+ debug: bool
+ Whether print debug logs.
+
+ Returns:
+ (proc, socket):
+ Yields a tuple with the subprocess as the first element and the UNIX-domain
+ IPC socket as the second element.
+ """
+
+ if not vineyardd_path:
+ vineyardd_path = find_vineyardd_path()
+
+ if not vineyardd_path:
+ raise RuntimeError('Unable to find the "vineyardd" executable')
+
+ if socket is None:
+ socketfp = tempfile.NamedTemporaryFile(
+ delete=True, prefix='vineyard-', suffix='.sock'
+ )
+ socket = socketfp.name
+ socketfp.close()
+
+ command = [
+ vineyardd_path,
+ '--deployment',
+ 'local',
+ '--size',
+ str(size),
+ '--socket',
+ socket,
+ '--rpc' if rpc else '--norpc',
+ '--rpc_socket_port',
+ str(rpc_socket_port),
+ '--meta',
+ meta,
+ ]
+
+ if meta == 'etcd':
+ if etcd_endpoints is None:
+ etcd_ctx = start_etcd()
+ (
+ _etcd_proc,
+ etcd_endpoints,
+ ) = etcd_ctx.__enter__() # pylint: disable=no-member
+ else:
+ etcd_ctx = None
+ command.extend(('--etcd_endpoint', etcd_endpoints))
+ if etcd_prefix is not None:
+ command.extend(('--etcd_prefix', etcd_prefix))
+ else:
+ etcd_ctx = None
+
+ env = os.environ.copy()
+ if debug:
+ env['GLOG_v'] = '11'
+
+ proc = None
+ try:
+ proc = subprocess.Popen(
+ command,
+ env=env,
+ stdout=subprocess.PIPE,
+ stderr=sys.__stderr__,
+ universal_newlines=True,
+ encoding='utf-8',
+ )
+ # wait for vineyardd ready: check the rpc port and ipc sockets
+ rc = proc.poll()
+ while rc is None:
+ if check_socket(socket) and (
+ (not rpc) or check_socket(('0.0.0.0', rpc_socket_port))
+ ):
+ break
+ time.sleep(1)
+ rc = proc.poll()
+
+ if rc is not None:
+ err = textwrap.indent(proc.stdout.read(), ' ' * 4)
+ raise RuntimeError(
+ 'vineyardd exited unexpectedly '
+ 'with code %d, error is:\n%s' % (rc, err)
+ )
+
+ logger.debug('vineyardd is ready.............')
+ yield proc, socket, etcd_endpoints
+ finally:
+ logger.debug('Local vineyardd being killed')
+ if proc is not None and proc.poll() is None:
+ proc.terminate()
+ proc.wait()
+ try:
+ shutil.rmtree(socket)
+ except Exception: # pylint: disable=broad-except
+ pass
+ if etcd_ctx is not None:
+ etcd_ctx.__exit__(None, None, None) # pylint: disable=no-member
+
+
+__default_instance_contexts = {}
+
+
+def try_init() -> str:
+ """
+ Launching a local vineyardd instance and get a client as easy as possible.
+
+ In a clean environment, simply use:
+
+ .. code:: python
+
+ vineyard.try_init()
+
+ It will launch a local vineyardd and return a connected client to the
+ vineyardd.
+
+ It will also setup the environment variable :code:`VINEYARD_IPC_SOCKET`.
+
+ The init method can only be called once in a process, to get the established
+ sockets or clients later in the process, use :code:`get_current_socket` or
+ :code:`connect()` respectively.
+ """
+ assert not __default_instance_contexts
+
+ if 'VINEYARD_IPC_SOCKET' in os.environ:
+ raise ValueError(
+ "VINEYARD_IPC_SOCKET has already been set: %s, which "
+ "means there might be a vineyard daemon already running "
+ "locally" % os.environ['VINEYARD_IPC_SOCKET']
+ )
+
+ ctx = start_vineyardd(meta='local', rpc=False)
+ _, ipc_socket, etcd_endpoints = ctx.__enter__() # pylint: disable=no-member
+ __default_instance_contexts[ipc_socket] = ctx
+
+ # populate the environment variable
+ os.environ['VINEYARD_IPC_SOCKET'] = ipc_socket
+ return get_current_socket()
+
+
+[docs]def get_current_socket() -> str:
+ """
+ Get current vineyard UNIX-domain socket established by :code:`vineyard.try_init()`.
+
+ Raises:
+ ValueError if vineyard is not initialized.
+ """
+ if not __default_instance_contexts:
+ raise ValueError(
+ 'Vineyard has not been initialized, '
+ 'use vineyard.try_init() to launch vineyard instances'
+ )
+ sockets = list(__default_instance_contexts.keys())
+ if sockets:
+ return sockets[0]
+ raise RuntimeError("Vineyard has not been initialized")
+
+
+def shutdown():
+ """
+ Shutdown the vineyardd instances launched by previous :code:`vineyard.try_init()`.
+ """
+ global __default_instance_contexts
+ if __default_instance_contexts:
+ for ipc_socket in reversed(__default_instance_contexts.keys()):
+ __default_instance_contexts[ipc_socket].__exit__(None, None, None)
+ # NB. don't pop pre-existing env if we not launch
+ os.environ.pop('VINEYARD_IPC_SOCKET', None)
+ __default_instance_contexts = {}
+
+
+@atexit.register
+def __shutdown_handler():
+ with contextlib.suppress(Exception):
+ shutdown()
+
+
+__all__ = ['start_vineyardd', 'try_init', 'shutdown']
+
+#! /usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2020-2023 Alibaba Group Holding Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+''' This module exposes support for ByteStream, that use can used like:
+
+.. code:: python
+
+ # create a builder, then seal it as stream
+ >>> stream = ByteStream.new(client)
+ >>> stream
+ ByteStream <o0001e09ddd98fd70>
+
+ # use write to put chunks
+ >>> writer = stream.open_writer(client)
+ >>> chunk = reader.next(1024)
+ >>> chunk
+ <memory at 0x136ca2ac0>
+ >>> len(chunk)
+ 1024
+ >>> chunk.readonly
+ False
+ >>> vineyard.memory_copy(chunk, src=b'abcde', offset=0)
+
+ # mark the stream as finished
+ >>> writer.finish()
+
+ # open a reader
+ >>> reader = stream.open_reader(client)
+ >>> chunk = reader.next()
+ >>> chunk
+ <memory at 0x136d207c0>
+ >>> len(chunk)
+ 1234
+ >>> chunk.readonly
+ True
+ >>> bytes(chunk[:10])
+ b'abcde\x00\x00\x00\x00\x00'
+
+ # the reader reaches the end of the stream
+ >>> chunk = reader.next()
+ ---------------------------------------------------------------------------
+ StreamDrainedException Traceback (most recent call last)
+ ~/libvineyard/python/vineyard/io/byte.py in next(self)
+ 108
+ --> 109 def next(self) -> memoryview:
+ 110 try:
+
+ StreamDrainedException: Stream drain: Stream drained: no more chunks
+
+ The above exception was the direct cause of the following exception:
+
+ StopIteration Traceback (most recent call last)
+ <ipython-input-11-d8809de11870> in <module>
+ ----> 1 chunk = reader.next()
+
+ ~/libvineyard/python/vineyard/io/byte.py in next(self)
+ 109 def next(self) -> memoryview:
+ 110 try:
+ --> 111 return self._client.next_buffer_chunk(self._stream)
+ 112 except StreamDrainedException as e:
+ 113 raise StopIteration('No more chunks') from e
+
+ StopIteration: No more chunks
+'''
+
+import json
+from io import BytesIO
+from typing import Dict
+
+from vineyard._C import ObjectID
+from vineyard._C import ObjectMeta
+from vineyard._C import StreamDrainedException
+from vineyard._C import memory_copy
+from vineyard.io.stream import BaseStream
+
+
+[docs]class ByteStream(BaseStream):
+ def __init__(self, meta: ObjectMeta, params: Dict = None):
+ super().__init__(meta)
+ self._params = params
+
+ @property
+ def params(self):
+ return self._params
+
+ @staticmethod
+ def new(client, params: Dict = None, meta: ObjectMeta = None) -> "ByteStream":
+ if meta is None:
+ meta = ObjectMeta()
+ meta['typename'] = 'vineyard::ByteStream'
+ if params is None:
+ params = dict()
+ meta['params_'] = params
+ meta = client.create_metadata(meta)
+ client.create_stream(meta.id)
+ return ByteStream(meta, params)
+
+ class Reader(BaseStream.Reader):
+ def __init__(self, client, stream: ObjectID):
+ super().__init__(client, stream)
+
+ def next(self) -> memoryview:
+ try:
+ return self._client.next_buffer_chunk(self._stream)
+ except StreamDrainedException as e:
+ raise StopIteration('No more chunks') from e
+
+ class Writer(BaseStream.Writer):
+ def __init__(self, client, stream: ObjectID):
+ super().__init__(client, stream)
+
+ self._buffer_size_limit = 1024 * 1024 * 256
+ self._buffer = BytesIO()
+
+ @property
+ def buffer_size_limit(self):
+ return self._buffer_size_limit
+
+ @buffer_size_limit.setter
+ def buffer_size_limit(self, value: int):
+ self._buffer_size_limit = value
+
+ def next(self, size: int) -> memoryview:
+ return self._client.new_buffer_chunk(self._stream, size)
+
+ def write(self, data: bytes):
+ self._buffer.write(data)
+ self._try_flush_buffer()
+
+ def _try_flush_buffer(self, force=False):
+ view = self._buffer.getbuffer()
+ if len(view) >= self._buffer_size_limit or (force and len(view) > 0):
+ if len(view) > 0:
+ chunk = self.next(len(view))
+ memory_copy(chunk, view)
+ self._buffer = BytesIO()
+
+ def finish(self):
+ self._try_flush_buffer(True)
+ return self._client.stop_stream(self._stream, False)
+
+ def _open_new_reader(self, client):
+ return ByteStream.Reader(client, self.id)
+
+ def _open_new_writer(self, client):
+ return ByteStream.Writer(client, self.id)
+
+
+def byte_stream_resolver(obj):
+ meta = obj.meta
+ if 'params_' in meta:
+ params = json.loads(meta['params_'])
+ else:
+ params = dict
+ return ByteStream(obj.meta, params)
+
+
+def register_byte_stream_types(_builder_ctx, resolver_ctx):
+ if resolver_ctx is not None:
+ resolver_ctx.register('vineyard::ByteStream', byte_stream_resolver)
+
+#! /usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2020-2023 Alibaba Group Holding Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+''' This module exposes support for DataframeStream, that use can used like:
+
+.. code:: python
+
+ # create a builder, then seal it as stream
+ >>> stream = DataframeStream.new(client)
+ >>> stream = builder.seal(client)
+ >>> stream
+ DataframeStream <o0001e09ddd98fd70>
+
+ # use write to put chunks
+ >>> writer = stream.open_writer(client)
+ >>> writer.write_table(
+ pa.Table.from_pandas(
+ pd.DataFrame({"x": [1,2,3], "y": [4,5,6]})))
+
+ # mark the stream as finished
+ >>> writer.finish()
+
+ # open a reader
+ >>> reader = stream.open_reader(client)
+ >>> batch = reader.next()
+ >>> batch
+ pyarrow.RecordBatch
+ x: int64
+ y: int64
+
+ # the reader reaches the end of the stream
+ >>> batch = reader.next()
+ ---------------------------------------------------------------------------
+ StreamDrainedException Traceback (most recent call last)
+ ~/libvineyard/python/vineyard/io/dataframe.py in next(self)
+ 97 try:
+ ---> 98 buffer = self._client.next_buffer_chunk(self._stream)
+ 99 with pa.ipc.open_stream(buffer) as reader:
+
+ StreamDrainedException: Stream drain: Stream drained: no more chunks
+
+ The above exception was the direct cause of the following exception:
+
+ StopIteration Traceback (most recent call last)
+ <ipython-input-11-10f09bf65f8a> in <module>
+ ----> 1 batch = reader.next()
+
+ ~/libvineyard/python/vineyard/io/dataframe.py in next(self)
+ 100 return reader.read_next_batch()
+ 101 except StreamDrainedException as e:
+ --> 102 raise StopIteration('No more chunks') from e
+ 103
+ 104 def __str__(self) -> str:
+
+ StopIteration: No more chunks
+'''
+
+import contextlib
+import json
+from io import BytesIO
+from typing import Dict
+
+import pyarrow as pa
+import pyarrow.ipc # pylint: disable=unused-import
+
+from vineyard._C import ObjectID
+from vineyard._C import ObjectMeta
+from vineyard._C import StreamDrainedException
+from vineyard._C import memory_copy
+from vineyard.io.stream import BaseStream
+
+
+[docs]class DataframeStream(BaseStream):
+ def __init__(self, meta: ObjectMeta, params: Dict = None):
+ super().__init__(meta)
+ self._params = params
+
+ @property
+ def params(self):
+ return self._params
+
+ @staticmethod
+ def new(client, params: Dict = None, meta: ObjectMeta = None) -> "DataframeStream":
+ if meta is None:
+ meta = ObjectMeta()
+ meta['typename'] = 'vineyard::DataframeStream'
+ if params is None:
+ params = dict()
+ meta['params_'] = params
+ meta = client.create_metadata(meta)
+ client.create_stream(meta.id)
+ return DataframeStream(meta, params)
+
+ class Reader(BaseStream.Reader):
+ def __init__(self, client, stream: ObjectID):
+ super().__init__(client, stream)
+
+ def next(self) -> pa.RecordBatch:
+ try:
+ buffer = self._client.next_buffer_chunk(self._stream)
+ with pa.ipc.open_stream(buffer) as reader:
+ return reader.read_next_batch()
+ except StreamDrainedException as e:
+ raise StopIteration('No more chunks') from e
+
+ def read_table(self) -> pa.Table:
+ batches = []
+ while True:
+ try:
+ batches.append(self.next())
+ except StopIteration:
+ break
+ return pa.Table.from_batches(batches)
+
+ class Writer(BaseStream.Writer):
+ def __init__(self, client, stream: ObjectID):
+ super().__init__(client, stream)
+
+ self._buffer = BytesIO()
+
+ def next(self, size: int) -> memoryview:
+ return self._client.new_buffer_chunk(self._stream, size)
+
+ def write(self, batch: pa.RecordBatch):
+ sink = BytesIO()
+ with pa.ipc.new_stream(sink, batch.schema) as writer:
+ writer.write(batch)
+ view = sink.getbuffer()
+ if len(view) > 0:
+ buffer = self.next(len(view))
+ memory_copy(buffer, view)
+
+ def write_table(self, table: pa.Table):
+ for batch in table.to_batches():
+ self.write(batch)
+
+ def finish(self):
+ return self._client.stop_stream(self._stream, False)
+
+ def _open_new_reader(self, client):
+ return DataframeStream.Reader(client, self.id)
+
+ def _open_new_writer(self, client):
+ return DataframeStream.Writer(client, self.id)
+
+
+def dataframe_stream_resolver(obj):
+ meta = obj.meta
+ if 'params_' in meta:
+ params = json.loads(meta['params_'])
+ else:
+ params = dict
+ return DataframeStream(obj.meta, params)
+
+
+def register_dataframe_stream_types(_builder_ctx, resolver_ctx):
+ if resolver_ctx is not None:
+ resolver_ctx.register('vineyard::DataframeStream', dataframe_stream_resolver)
+
+#! /usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2020-2023 Alibaba Group Holding Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+''' This module exposes support for RecordBatchStream.
+'''
+
+import contextlib
+import json
+from typing import Any
+from typing import Dict
+from typing import Optional
+
+from vineyard._C import ObjectMeta
+from vineyard.core import context
+from vineyard.io.stream import BaseStream
+
+
+[docs]class RecordBatchStream(BaseStream):
+ def __init__(self, meta: ObjectMeta, params: Optional[Dict[str, Any]] = None):
+ super().__init__(meta)
+ self._params = params
+
+ @property
+ def params(self):
+ return self._params
+
+ @staticmethod
+ def new(
+ client,
+ params: Optional[Dict[str, Any]] = None,
+ meta: Optional[ObjectMeta] = None,
+ ) -> "RecordBatchStream":
+ if meta is None:
+ meta = ObjectMeta()
+ meta['typename'] = 'vineyard::RecordBatchStream'
+ if params is None:
+ params = dict()
+ meta['params_'] = params
+ meta = client.create_metadata(meta)
+ client.create_stream(meta.id)
+ return RecordBatchStream(meta, params)
+
+
+def recordbatch_stream_resolver(obj, resolver): # pylint: disable=unused-argument
+ meta = obj.meta
+ if 'params_' in meta:
+ params = json.loads(meta['params_'])
+ else:
+ params = dict
+ return RecordBatchStream(meta, params)
+
+
+def register_recordbatch_stream_types(_builder_ctx, resolver_ctx):
+ if resolver_ctx is not None:
+ resolver_ctx.register(
+ 'vineyard::RecordBatchStream', recordbatch_stream_resolver
+ )
+
+
+@contextlib.contextmanager
+def recordbatch_stream_context():
+ with context() as (builder_ctx, resolver_ctx):
+ register_recordbatch_stream_types(builder_ctx, resolver_ctx)
+ yield builder_ctx, resolver_ctx
+
+#! /usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2020-2023 Alibaba Group Holding Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import json
+import logging
+import traceback
+from typing import Callable
+from typing import Dict
+from typing import List
+from typing import Optional
+from urllib.parse import urlparse
+
+from vineyard._C import ObjectID
+from vineyard._C import ObjectMeta
+from vineyard._C import StreamDrainedException
+from vineyard.core.driver import registerize
+from vineyard.core.resolver import resolver_context
+
+logger = logging.getLogger('vineyard')
+
+
+[docs]@registerize
+def read(
+ path,
+ *args,
+ handlers=None,
+ accumulate=False,
+ chunk_hook: Optional[Callable] = None,
+ **kwargs
+):
+ """Open a path and read it as a single stream.
+
+ Parameters
+ ----------
+ path: str
+ Path to read, the last reader registered for the scheme of
+ the path will be used.
+ handlers: list, optional
+ If handlers is not None, launched worker processes will be
+ emplaced into the list for further customized job lifecycle
+ management. Default is None.
+ accumulate: bool, optional
+ If :code:`accumulate` is True, it will return a data frame,
+ rather than dataframe stream. Default is False.
+ chunk_hook: callable, optional
+ If the read/write target is a global dataframe (e.g., csv,
+ orc, parquet, etc.), the hook will be called for each chunk
+ to be read or write (usually a :code:`pyarrow.RecordBatch`).
+ The hook should return a :code:`pyarrow.RecordBatch` object
+ and should be stateless as the invoke order is not guaranteed.
+ E.g.,
+
+ .. code:: python
+
+ def exchange_column(batch):
+ import pyarrow as pa
+
+ columns = batch.columns
+ first = columns[0]
+ second = columns[1]
+ columns = [second, first] + columns[2:]
+ return pa.RecordBatch.from_arrays(columns, schema=batch.schema)
+
+ vineyard_ipc_socket: str
+ The local or remote vineyard's IPC socket location that the
+ remote readers will use to establish connections with the
+ vineyard server.
+ vineyard_endpoint: str, optional
+ An optional address of vineyard's RPC socket, which will be
+ used for retrieving server's information on the client side.
+ If not provided, the `vineyard_ipc_socket` will be used, or
+ it will tries to discovery vineyard's IPC or RPC endpoints
+ from environment variables.
+ """
+ parsed = urlparse(path)
+ if read._factory and read._factory.get(parsed.scheme):
+ errors = []
+ for reader in read._factory[parsed.scheme][::-1]:
+ try:
+ proc_kwargs = kwargs.copy()
+ r = reader(
+ path,
+ proc_kwargs.pop('vineyard_ipc_socket'),
+ *args,
+ handlers=handlers,
+ accumulate=accumulate,
+ chunk_hook=chunk_hook,
+ **proc_kwargs
+ )
+ if r is not None:
+ return r
+ except Exception: # pylint: disable=broad-except
+ errors.append('%s: %s' % (reader.__name__, traceback.format_exc()))
+ raise RuntimeError(
+ 'Unable to find a proper IO driver for %s, potential causes are:\n %s'
+ % (path, '\n'.join(errors))
+ )
+ else:
+ raise ValueError("No IO driver registered for %s" % path)
+
+
+[docs]@registerize
+def write(
+ path, stream, *args, handlers=None, chunk_hook: Optional[Callable] = None, **kwargs
+):
+ """Write the stream to a given path.
+
+ Parameters
+ ----------
+ path: str
+ Path to write, the last writer registered for the scheme of the path
+ will be used.
+ stream: vineyard stream
+ Stream that produces the data to write.
+ handlers: list, optional
+ If handlers is not None, launched worker processes will be
+ emplaced into the list for further customized job lifecycle
+ management. Default is None.
+ chunk_hook: callable, optional
+ If the read/write target is a global dataframe (e.g., csv,
+ orc, parquet, etc.), the hook will be called for each chunk
+ to be read or write (usually a :code:`pyarrow.RecordBatch`).
+ The hook should return a :code:`pyarrow.RecordBatch` object
+ and should be stateless as the invoke order is not guaranteed.
+ E.g.,
+
+ .. code:: python
+
+ def exchange_column(batch):
+ import pyarrow as pa
+
+ columns = batch.columns
+ first = columns[0]
+ second = columns[1]
+ columns = [second, first] + columns[2:]
+ return pa.RecordBatch.from_arrays(columns, schema=batch.schema)
+
+ vineyard_ipc_socket: str
+ The local or remote vineyard's IPC socket location that the remote
+ readers will use to establish connections with the vineyard server.
+ vineyard_endpoint: str, optional
+ An optional address of vineyard's RPC socket, which will be used for
+ retrieving server's information on the client side. If not provided,
+ the `vineyard_ipc_socket` will be used, or it will tries to discovery
+ vineyard's IPC or RPC endpoints from environment variables.
+ """
+ parsed = urlparse(path)
+ if write._factory and write._factory.get(parsed.scheme):
+ errors = []
+ for writer in write._factory[parsed.scheme][::-1]:
+ try:
+ proc_kwargs = kwargs.copy()
+ writer(
+ path,
+ stream,
+ proc_kwargs.pop('vineyard_ipc_socket'),
+ *args,
+ handlers=handlers,
+ chunk_hook=chunk_hook,
+ **proc_kwargs
+ )
+ except Exception: # pylint: disable=broad-except
+ exc = traceback.format_exc()
+ errors.append('%s: %s' % (writer.__name__, exc))
+ if 'StreamFailedException: Stream failed' in exc:
+ # if the source stream has already failed, we should
+ # fail immediately and don't try other drivers.
+ break
+ continue
+ else:
+ return
+ raise RuntimeError(
+ 'Unable to find a proper IO driver for %s, potential causes are:\n %s'
+ % (path, '\n'.join(errors))
+ )
+ else:
+ raise ValueError("No IO driver registered for %s" % path)
+
+
+[docs]def open(
+ path,
+ *args,
+ mode='r',
+ handlers=None,
+ chunk_hook: Optional[Callable] = None,
+ **kwargs
+):
+ """Open a path as a reader or writer, depends on the parameter :code:`mode`.
+ If :code:`mode` is :code:`r`, it will open a stream for read, and open a
+ stream for write when :code:`mode` is :code:`w`.
+
+ Parameters
+ ----------
+ path: str
+ Path to open.
+ mode: char
+ Mode about how to open the path, :code:`r` is for read and :code:`w` for write.
+ handlers:
+ A dict that will be filled with a :code:`handler` that contains the process
+ handler of the underlying read/write process that can be joined using
+ :code:`join` to capture the possible errors during the I/O proceeding.
+ chunk_hook: callable, optional
+ If the read/write target is a global dataframe (e.g., csv,
+ orc, parquet, etc.), the hook will be called for each chunk
+ to be read or write (usually a :code:`pyarrow.RecordBatch`).
+ The hook should return a :code:`pyarrow.RecordBatch` object
+ and should be stateless as the invoke order is not guaranteed.
+ E.g.,
+
+ .. code:: python
+
+ def exchange_column(batch):
+ import pyarrow as pa
+
+ columns = batch.columns
+ first = columns[0]
+ second = columns[1]
+ columns = [second, first] + columns[2:]
+ return pa.RecordBatch.from_arrays(columns, schema=batch.schema)
+
+ vineyard_ipc_socket: str
+ Vineyard's IPC socket location.
+ vineyard_endpoint: str
+ Vineyard's RPC socket address.
+
+ See Also
+ --------
+ vineyard.io.read
+ vineyard.io.write
+ """
+ parsed = urlparse(path)
+ if not parsed.scheme:
+ path = 'file://' + path
+
+ if mode == 'r':
+ return read(path, *args, handlers=handlers, chunk_hook=chunk_hook, **kwargs)
+
+ if mode == 'w':
+ return write(path, *args, handlers=handlers, chunk_hook=chunk_hook, **kwargs)
+
+ raise RuntimeError('Opening %s with mode %s is not supported' % (path, mode))
+
+
+class BaseStream:
+ class Reader:
+ def __init__(self, client, stream: ObjectID, resolver=None):
+ self._client = client
+ self._stream = stream
+ self._resolver = resolver
+ self._client.open_stream(stream, 'r')
+
+ def next(self) -> object:
+ try:
+ chunk = self._client.next_chunk(self._stream)
+ except StreamDrainedException as e:
+ raise StopIteration('No more chunks') from e
+
+ if self._resolver is not None:
+ return self._resolver(chunk)
+ else:
+ with resolver_context() as ctx:
+ return ctx.run(chunk)
+
+ def next_metadata(self) -> ObjectMeta:
+ try:
+ return self._client.next_chunk_meta(self._stream)
+ except StreamDrainedException as e:
+ raise StopIteration('No more chunks') from e
+
+ def __str__(self) -> str:
+ return repr(self)
+
+ def __repr__(self) -> str:
+ return '%s of Stream <%r>' % (self.__class__, self._stream)
+
+ class Writer:
+ def __init__(self, client, stream: ObjectID):
+ self._client = client
+ self._stream = stream
+ self._client.open_stream(stream, 'w')
+
+ def next(self, size: int) -> memoryview:
+ return self._client.new_buffer_chunk(self._stream, size)
+
+ def append(self, chunk: ObjectID):
+ return self._client.push_chunk(self._stream, chunk)
+
+ def fail(self):
+ return self._client.stop_stream(self._stream, True)
+
+ def finish(self):
+ return self._client.stop_stream(self._stream, False)
+
+ def __str__(self) -> str:
+ return repr(self)
+
+ def __repr__(self) -> str:
+ return '%s of Stream <%r>' % (self.__class__, self._stream)
+
+ def __init__(self, meta: ObjectMeta, resolver=None):
+ self._meta = meta
+ self._stream = meta.id
+ self._resolver = resolver
+
+ self._reader = None
+ self._writer = None
+
+ @property
+ def id(self) -> ObjectID:
+ return self._stream
+
+ @property
+ def meta(self) -> ObjectMeta:
+ return self._meta
+
+ @property
+ def reader(self) -> "BaseStream.Reader":
+ return self.open_reader()
+
+ def __str__(self) -> str:
+ return repr(self)
+
+ def __repr__(self) -> str:
+ return '%s <%r>' % (self.__class__.__name__, self._stream)
+
+ def _open_new_reader(self, client) -> "BaseStream.Reader":
+ '''Always open a new reader.'''
+ return BaseStream.Reader(client, self.id, self._resolver)
+
+ def open_reader(self, client=None) -> "BaseStream.Reader":
+ if self._reader is None:
+ if client is None:
+ client = self._meta._client
+ self._reader = self._open_new_reader(client)
+ return self._reader
+
+ @property
+ def writer(self) -> "BaseStream.Writer":
+ return self.open_writer()
+
+ def _open_new_writer(self, client) -> "BaseStream.Writer":
+ return BaseStream.Writer(client, self.id)
+
+ def open_writer(self, client=None) -> "BaseStream.Writer":
+ if self._writer is None:
+ if client is None:
+ client = self._meta._client
+ self._writer = self._open_new_writer(client)
+ return self._writer
+
+ def drop(self, client=None):
+ if client is None:
+ client = self._meta._client
+ if hasattr(client, 'drop_stream'):
+ client.drop_stream(self.id)
+
+
+class StreamCollection:
+ """A stream collection is a set of stream, where each element is a stream, or,
+ another stream collection.
+ """
+
+ KEY_OF_STREAMS = '__streams'
+ KEY_OF_PATH = '__path'
+ KEY_OF_GLOBAL = '__global'
+ KEY_OF_OPTIONS = '__options'
+
+ def __init__(self, meta: ObjectMeta, streams: List[ObjectID]):
+ self._meta = meta
+ self._streams = streams
+ if StreamCollection.KEY_OF_GLOBAL in self._meta:
+ self._global = self._meta[StreamCollection.KEY_OF_GLOBAL]
+ else:
+ self._global = False
+
+ @staticmethod
+ def new(
+ client, metadata: Dict, streams: List[ObjectID], meta: ObjectMeta = None
+ ) -> "StreamCollection":
+ if meta is None:
+ meta = ObjectMeta()
+ meta['typename'] = 'vineyard::StreamCollection'
+ for k, v in metadata.items():
+ if k not in [
+ 'id',
+ 'signature',
+ 'instance_id',
+ 'transient',
+ 'global',
+ 'typename',
+ ]:
+ meta[k] = v
+ meta[StreamCollection.KEY_OF_STREAMS] = [int(s) for s in streams]
+ meta = client.create_metadata(meta)
+ return StreamCollection(meta, streams)
+
+ @property
+ def id(self):
+ return self.meta.id
+
+ @property
+ def meta(self):
+ return self._meta
+
+ @property
+ def isglobal(self):
+ return self._global
+
+ @property
+ def streams(self):
+ return self._streams
+
+ def __repr__(self) -> str:
+ return "StreamCollection: %s [%s]" % (
+ repr(self.id),
+ [repr(s) for s in self.streams],
+ )
+
+ def __str__(self) -> str:
+ return repr(self)
+
+
+def stream_collection_resolver(obj):
+ meta = obj.meta
+ streams = json.loads(meta[StreamCollection.KEY_OF_STREAMS])
+ return StreamCollection(meta, [ObjectID(s) for s in streams])
+
+
+def register_stream_collection_types(_builder_ctx, resolver_ctx):
+ if resolver_ctx is not None:
+ resolver_ctx.register('vineyard::StreamCollection', stream_collection_resolver)
+
+
+__all__ = [
+ 'open',
+ 'read',
+ 'write',
+ 'BaseStream',
+ 'StreamCollection',
+ 'register_stream_collection_types',
+]
+
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+#
+# Copyright 2020-2023 Alibaba Group Holding Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+'''
+The shared_memory module provides similar interface like multiprocessing.shared_memory
+for direct access shared memory backed by vineyard across processes.
+
+The API is kept consistent with multiprocessing.shared_memory but the semantics is
+slightly different. For vineyard, to make the shared memory visible for other process,
+a explicitly ``seal`` or ``close`` operation is needed.
+
+Refer to the documentation of multiprocessing.shared_memory for details.
+'''
+
+# pylint: skip-file
+
+try:
+ import multiprocessing.shared_memory as shm
+except ImportError:
+ # multiprocessing.shared_memory is available since Python 3.8, we use the slim
+ # library for earlier version of Python.
+ #
+ # see also github #327.
+ import shared_memory as shm
+
+import struct
+import warnings
+
+from vineyard._C import ObjectID
+
+
+
+
+
+_encoding = "utf8"
+
+
+
+
+
+__all__ = ['SharedMemory', 'ShareableList']
+