diff --git a/python/vineyard/core/client.py b/python/vineyard/core/client.py index cdd1611c9d..86b035c759 100644 --- a/python/vineyard/core/client.py +++ b/python/vineyard/core/client.py @@ -34,6 +34,7 @@ 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 @@ -49,6 +50,44 @@ def _apply(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.yaml`. + + 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 = os.path.join(config, 'vineyard.yaml') + 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 and 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) + return ipc_socket, rpc_endpoint + + 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. @@ -64,14 +103,53 @@ def __init__( session: int = None, username: str = None, password: str = None, + config: str = None, ): """Connects to the vineyard IPC socket and RPC socket. - - The argument `socket` takes precedence over environment variable - `VINEYARD_IPC_SOCKET` for IPC client. - - The argument `endpoint` takes precedence over arguments (`host`, `port`), - which further takes precedence over environment variable - `VINEYARD_RPC_ENDPOINT` for RPC client. + - 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.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.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 @@ -96,16 +174,36 @@ def __init__( endpoint = endpoint.split(':') host, port = endpoint + hosts, ports = [], [] + 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) - if host and port: - self._rpc_client = _connect(host, port, **kwargs) + for host, port in zip(hosts, ports): + try: + self._rpc_client = _connect(host, port, **kwargs) + break + except VineyardException: + continue 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 are not and environment variables VINEYARD_IPC_SOCKET " - "and VINEYARD_RPC_ENDPOINT are not available." + "Arguments, environment variables `VINEYARD_IPC_SOCKET` " + "and `VINEYARD_RPC_ENDPOINT`, as well as the configuration file, " + "are all unavailable." ) @property diff --git a/requirements.txt b/requirements.txt index c2d16692b6..683ae641d2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,6 +10,7 @@ pandas>=1.0.0; python_version>="3.7" pickle5; python_version<="3.7" psutil pyarrow +pyyaml setuptools shared-memory38; python_version<="3.7" sortedcontainers