Add the config parameter to Client to support loading configuration…

… from YAML. (#1689)

Signed-off-by: Ye Cao <[email protected]>
Signed-off-by: Tao He <[email protected]>
Co-authored-by: Tao He <[email protected]>

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

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