config: update bento config for yatai server (#1782)

* update bento config for yatai server

* format

* add yatai client utils

* update

* update simple-di version

Author: yubozhao

bentoml/__init__.py

@@ -4,23 +4,24 @@
 __version__ = get_versions()["version"]  # type: ignore
 del get_versions
 
-# from bentoml._internal.configuration import inject_dependencies
+from bentoml._internal.configuration import inject_dependencies
+
 # from bentoml._internal.utils.log import configure_logging
-#
-# # Inject dependencies and configurations
-# inject_dependencies()
+
+# Inject dependencies and configurations
+inject_dependencies()
 #
 # # Configuring logging properly before loading other modules
 # configure_logging()
 
-from ._internal.bundle import containerize, load
-from ._internal.environment import env
-from ._internal.inference_api import api, batch_api
-from ._internal.models import Model, PickleModel
-from ._internal.repository import delete, get, ls, pull, push
-from ._internal.server import serve
-from ._internal.service import Service
-from ._internal.yatai_client import YataiClient
+from bentoml._internal.bundle import containerize, load
+from bentoml._internal.environment import env
+from bentoml._internal.inference_api import api, batch_api
+from bentoml._internal.models import Model, PickleModel
+from bentoml._internal.repository import delete, get, ls, pull, push
+from bentoml._internal.server import serve
+from bentoml._internal.service import Service
+from bentoml._internal.yatai_client import YataiClient
 
 __all__ = [
     "__version__",

bentoml/_internal/configuration/containers.py

@@ -8,12 +8,11 @@
 from simple_di import Provide, Provider, container, providers
 
 from bentoml import __version__
+from bentoml._internal.configuration import expand_env_var, get_bentoml_deploy_version
+from bentoml._internal.utils import get_free_port
+from bentoml._internal.utils.ruamel_yaml import YAML
 from bentoml.exceptions import BentoMLConfigException
 
-from ..configuration import expand_env_var, get_bentoml_deploy_version
-from ..utils import get_free_port
-from ..utils.ruamel_yaml import YAML
-
 if TYPE_CHECKING:
     from ..marshal.marshal import MarshalApp
 
@@ -79,36 +78,20 @@
         },
         "adapters": {"image_input": {"default_extensions": [str]}},
         "yatai": {
-            "remote": {
-                "url": Or(str, None),
-                "access_token": Or(str, None),
-                "access_token_header": Or(str, None),
-                "tls": {
-                    "root_ca_cert": Or(str, None),
-                    "client_key": Or(str, None),
-                    "client_cert": Or(str, None),
-                    "client_certificate_file": Or(str, None),
-                },
-            },
-            "repository": {
-                "type": And(
-                    str,
-                    lambda _type: _type in YATAI_REPOSITORY_TYPES,
-                    error="yatai.repository.type must be one of %s"
-                    % YATAI_REPOSITORY_TYPES,
-                ),
-                "file_system": {"directory": Or(str, None)},
-                "s3": {
+            "default_server": Or(str, None),
+            "servers": {
+                str: {
                     "url": Or(str, None),
-                    "endpoint_url": Or(str, None),
-                    "signature_version": Or(str, None),
-                    "expiration": Or(int, None),
+                    "access_token": Or(str, None),
+                    "access_token_header": Or(str, None),
+                    "tls": {
+                        "root_ca_cert": Or(str, None),
+                        "client_key": Or(str, None),
+                        "client_cert": Or(str, None),
+                        "client_certificate_file": Or(str, None),
+                    },
                 },
-                "gcs": {"url": Or(str, None), "expiration": Or(int, None)},
             },
-            "database": {"url": Or(str, None)},
-            "namespace": str,
-            "logging": {"path": Or(str, None)},
         },
     }
 )
@@ -313,13 +296,6 @@ def metrics_client(
             namespace=namespace,
         )
 
-    @providers.SingletonFactory
-    @staticmethod
-    def yatai_metrics_client():
-        from ..metrics.prometheus import PrometheusClient
-
-        return PrometheusClient(multiproc=False, namespace="YATAI")
-
     bento_bundle_deployment_version = providers.Factory(
         get_bentoml_deploy_version,
         providers.Factory(
@@ -328,26 +304,12 @@ def yatai_metrics_client():
             config.bento_bundle.deployment_version,
         ),
     )
+    default_yatai_server = config.yatai.servers[config.yatai.default_server]
 
-    yatai_database_url = providers.Factory(
-        lambda default, customized: customized or default,
-        providers.Factory(
-            "sqlite:///{}".format,
-            providers.Factory(os.path.join, bentoml_home, "storage.db"),
-        ),
-        config.yatai.database.url,
-    )
-
-    yatai_file_system_directory = providers.Factory(
-        lambda default, customized: customized or default,
-        providers.Factory(os.path.join, bentoml_home, "repository"),
-        config.yatai.repository.file_system.directory,
-    )
-
-    yatai_tls_root_ca_cert = providers.Factory(
+    default_yatai_server_tls_root_ca_cert = providers.Factory(
         lambda current, deprecated: current or deprecated,
-        config.yatai.remote.tls.root_ca_cert,
-        config.yatai.remote.tls.client_certificate_file,
+        config.yatai.servers[config.yatai.default_server].tls.root_ca_cert,
+        config.yatai.servers[config.yatai.default_server].tls.client_certificate_file,
     )
 
     logging_file_directory = providers.Factory(
@@ -356,11 +318,5 @@ def yatai_metrics_client():
         config.logging.file.directory,
     )
 
-    yatai_logging_path = providers.Factory(
-        lambda default, customized: customized or default,
-        providers.Factory(os.path.join, logging_file_directory, "yatai_web_server.log"),
-        config.yatai.logging.path,
-    )
-
 
 BentoMLContainer = BentoMLContainerClass()

bentoml/_internal/configuration/default_configuration.yml

@@ -61,29 +61,14 @@ adapters:
       - .bmp
 
 yatai:
-  remote:
-    url: Null
-    access_token: Null
-    access_token_header: access_token
-    tls:
-      root_ca_cert: Null
-      client_certificate_file: Null
-      client_key: Null
-      client_cert: Null
-  repository:
-    type: file_system
-    file_system:
-      directory: Null
-    s3:
-      url: Null
-      endpoint_url: Null
-      signature_version: s3v4
-      expiration: 3000
-    gcs:
+  default_server: default
+  servers:
+    default:
       url: Null
-      expiration: 3000
-  database:
-    url: Null
-  namespace: dev
-  logging:
-    path: Null
+      access_token: Null
+      access_token_header: access_token
+      tls:
+        root_ca_cert: Null
+        client_certificate_file: Null
+        client_key: Null
+        client_cert: Null

bentoml/_internal/yatai_client/__init__.py

@@ -10,12 +10,36 @@
 # logger = logging.getLogger(__name__)
 
 
-# TODO:
+import logging
+
+from simple_di import inject
+
+from bentoml._internal.utils import cached_property
+from bentoml._internal.yatai_client.bento_repository_api import BentoRepositoryAPIClient
+from bentoml._internal.yatai_client.deployment_api import DeploymentAPIClient
+from bentoml._internal.yatai_client.utils import parse_grpc_url
+
+logger = logging.getLogger(__name__)
+
+
 class YataiClient:
     """
     Python Client for interacting with YataiService
     """
 
+    def __init__(self, yatai_server_name: str):
+        self._yatai_service = get_yatai_service()
+        self.bundle_api_client = None
+        self.deploy_api_client = None
+
+    @cached_property
+    def bundles(self):
+        return BentoRepositoryAPIClient(self._yatai_service)
+
+    @cached_property
+    def deployment(self):
+        return DeploymentAPIClient(self._yatai_service)
+
     # def __init__(self, yatai_service: Optional["YataiStub"] = None):
     #     self.yatai_service = yatai_service if yatai_service else get_yatai_service()
     #     self.bento_repository_api_client = None
@@ -46,3 +70,45 @@ def get_yatai_client(yatai_url: str = None) -> "YataiClient":
 
     # yatai_service = get_yatai_service(channel_address=yatai_url)
     # return YataiClient(yatai_service=yatai_service)
+
+
+@inject
+def get_yatai_service(
+    channel_address: str,
+    access_token: str,
+    access_token_header: str,
+    tls_root_ca_cert: str,
+    tls_client_key: str,
+    tls_client_cert: str,
+):
+    import certifi
+    import grpc
+
+    from bentoml._internal.yatai_client.interceptor import header_client_interceptor
+    from bentoml._internal.yatai_client.proto.yatai_service_pb2_grpc import YataiStub
+
+    channel_address = channel_address.strip()
+    schema, addr = parse_grpc_url(channel_address)
+    header_adder_interceptor = header_client_interceptor.header_adder_interceptor(
+        access_token_header, access_token
+    )
+    if schema in ("grpc", "https"):
+        tls_root_ca_cert = tls_root_ca_cert or certifi.where()
+        with open(tls_client_cert, "rb") as fb:
+            ca_cert = fb.read()
+        if tls_client_key:
+            with open(tls_client_key, "rb") as fb:
+                tls_client_key = fb.read()
+        if tls_client_cert:
+            with open(tls_client_cert, "rb") as fb:
+                tls_client_cert = fb.read()
+        credentials = grpc.ssl_channel_credentials(
+            root_certificates=ca_cert,
+            private_key=tls_client_key,
+            certificate_chain=tls_client_cert,
+        )
+        channel = grpc.secure_channel(addr, credentials)
+    else:
+        channel = grpc.insecure_channel(addr)
+
+    return YataiStub(grpc.intercept_channel(channel, header_adder_interceptor))

bentoml/_internal/yatai_client/interceptor/__init__.py

@@ -0,0 +1,11 @@
+def parse_grpc_url(url):
+    """
+    >>> parse_grpc_url("grpcs://yatai.com:43/query")
+    ('grpcs', 'yatai.com:43/query')
+    >>> parse_grpc_url("yatai.com:43/query")
+    (None, 'yatai.com:43/query')
+    """
+    from urllib3.util import parse_url
+
+    parts = parse_url(url)
+    return parts.scheme, url.replace(f"{parts.scheme}://", "", 1)

yatai/yatai/interceptor/generic_client_interceptor.py renamed to bentoml/_internal/yatai_client/interceptor/generic_client_interceptor.py

@@ -155,4 +155,4 @@ def save(self, path: PathType) -> None:
         tokenizer_type = self._model.get("tokenizer").__class__.__name__
         self._model.get("model").save_pretrained(path)
         self._model.get("tokenizer").save_pretrained(path)
-        self.__save_model_type(path, tokenizer_type)
+        self.__save_model_type(path, tokenizer_type)

yatai/yatai/interceptor/header_client_interceptor.py renamed to bentoml/_internal/yatai_client/interceptor/header_client_interceptor.py

@@ -39,7 +39,7 @@
     'contextvars;python_version < "3.7"',
     'dataclasses;python_version < "3.7"',
     "chardet",
-    "simple-di==0.1.0",
+    "simple-di==0.1.1",
     "cloudpickle",
 ]
 

bentoml/_internal/yatai_client/utils.py

@@ -1,8 +1,6 @@
 from .grpc_server_interceptor import PromServerInterceptor, ServiceLatencyInterceptor
-from .header_client_interceptor import header_adder_interceptor
 
 __all__ = [
     "ServiceLatencyInterceptor",
     "PromServerInterceptor",
-    "header_adder_interceptor",
 ]