Improve service_map SDK parameter (#130)

Author: vhaldemar

src/yandex_cloud_ml_sdk/_auth.py

@@ -52,13 +52,11 @@ async def get_auth_metadata(
         """:meta private:"""
         # NB: we are can't create lock in Auth constructor, so we a reusing lock from client.
         # Look at client._lock doctstring for details.
-        pass
 
     @classmethod
     @abstractmethod
     async def applicable_from_env(cls, **_: Any) -> Self | None:
         """:meta private:"""
-        pass
 
 
 class NoAuth(BaseAuth):
@@ -402,7 +400,7 @@ async def _request_token(cls, timeout: float, metadata_url: str) -> str:
 async def get_auth_provider(
     *,
     auth: str | BaseAuth | None,
-    endpoint: str,
+    endpoint: str | None,
     yc_profile: str | None,
 ) -> BaseAuth:
     """

src/yandex_cloud_ml_sdk/_client.py

@@ -14,7 +14,7 @@
 from yandex.cloud.endpoint.api_endpoint_service_pb2_grpc import ApiEndpointServiceStub
 
 from ._auth import BaseAuth, get_auth_provider
-from ._exceptions import AioRpcError
+from ._exceptions import AioRpcError, UnknownEndpointError
 from ._retry import RETRY_KIND_METADATA_KEY, RetryKind, RetryPolicy
 from ._types.misc import PathLike, coerce_path
 from ._utils.lock import LazyLock
@@ -44,7 +44,7 @@ class AsyncCloudClient:
     def __init__(
         self,
         *,
-        endpoint: str,
+        endpoint: str | None,
         auth: BaseAuth | str | None,
         service_map: dict[str, str],
         interceptors: Sequence[grpc.aio.ClientInterceptor] | None,
@@ -78,6 +78,10 @@ def __init__(
 
     async def _init_service_map(self, timeout: float):
         metadata = await self._get_metadata(auth_required=False, timeout=timeout, retry_kind=RetryKind.SINGLE)
+
+        if not self._endpoint:
+            raise RuntimeError('This method should be never called while endpoint=None')
+
         channel = self._new_channel(self._endpoint)
         async with channel:
             stub = ApiEndpointServiceStub(channel)
@@ -89,8 +93,30 @@ async def _init_service_map(self, timeout: float):
             for endpoint in response.endpoints:
                 self._service_map[endpoint.id] = endpoint.address
 
+    async def _discover_service_endpoint(
+        self,
+        service_name: str,
+        stub_class: type[StubType],
+        timeout: float
+    ) -> str:
+        endpoint: str | None
         # TODO: add a validation for unknown services in override
-        self._service_map.update(self._service_map_override)
+        if endpoint := self._service_map_override.get(service_name):
+            return endpoint
+
+        if self._endpoint is None:
+            raise UnknownEndpointError(
+                "due to `endpoint` SDK param explicitly set to `None` you need to define "
+                f"{service_name!r} endpoint manually at `service_map` SDK param"
+            )
+
+        if not self._service_map:
+            await self._init_service_map(timeout=timeout)
+
+        if endpoint := self._service_map.get(service_name):
+            return endpoint
+
+        raise UnknownEndpointError(f'failed to find endpoint for {service_name=} and {stub_class=}')
 
     async def _get_metadata(
         self,
@@ -172,19 +198,12 @@ async def _get_channel(
                 return self._channels[stub_class]
 
             service_name = service_name if service_name else service_for_ctor(stub_class)
-            if not self._service_map:
-                await self._init_service_map(timeout=timeout)
-
-            if not (endpoint := self._service_map.get(service_name)):
-                # NB: this fix will work if service_map will change ai-assistant to ai-assistants
-                # (and retrospectively if user will stuck with this version)
-                # and if _service_for_ctor will change ai-assistants to ai-assistant
-                if service_name in ('ai-assistant', 'ai-assistants'):
-                    service_name = 'ai-assistant' if service_name == 'ai-assistants' else 'ai-assistants'
-                    if not (endpoint := self._service_map.get(service_name)):
-                        raise ValueError(f'failed to find endpoint for {service_name=} and {stub_class=}')
-                else:
-                    raise ValueError(f'failed to find endpoint for {service_name=} and {stub_class=}')
+
+            endpoint = await self._discover_service_endpoint(
+                service_name=service_name,
+                stub_class=stub_class,
+                timeout=timeout
+            )
 
             self._endpoints[stub_class] = endpoint
             channel = self._channels[stub_class] = self._new_channel(endpoint)

src/yandex_cloud_ml_sdk/_exceptions.py

@@ -18,6 +18,9 @@
 class YCloudMLError(Exception):
     pass
 
+class UnknownEndpointError(YCloudMLError):
+    pass
+
 
 class RunError(YCloudMLError):
     def __init__(self, code: int, message: str, details: list[Any] | None, operation_id: str):

src/yandex_cloud_ml_sdk/_sdk.py

@@ -60,7 +60,7 @@ def __init__(
         self,
         *,
         folder_id: str,
-        endpoint: UndefinedOr[str] = UNDEFINED,
+        endpoint: UndefinedOr[str] | None = UNDEFINED,
         auth: UndefinedOr[str | BaseAuth] = UNDEFINED,
         retry_policy: UndefinedOr[RetryPolicy] = UNDEFINED,
         yc_profile: UndefinedOr[str] = UNDEFINED,
@@ -76,13 +76,15 @@ def __init__(
         :type folder_id: str
         :param endpoint: domain:port pair for Yandex Cloud API or any other
             grpc compatible target.
+            In case of ``None`` passed it turns off service endpoint discovery mechanism
+            and requires ``service_map`` to be passed.
         :type endpoint: str
         :param auth: string with API Key, IAM token or one of yandex_cloud_ml_sdk.auth objects;
             in case of default Undefined value, there will be a mechanism to get token
             from environment
         :type api_key | BaseAuth: str
         :param service_map: a way to redefine endpoints for one or more cloud subservices
-            with a format of dict {service_name: service_address}.
+            with a format of dict ``{"service_name": "service_address"}``.
         :type service_map: Dict[str, str]
         :param enable_server_data_logging: when passed bool, we will add
             `x-data-logging-enabled: <value>` to all of requests, which will
@@ -145,7 +147,7 @@ def _init_domains(self) -> None:
                 resource = member(name=member_name, sdk=self)
                 setattr(self, member_name, resource)
 
-    def _get_endpoint(self, endpoint: UndefinedOr[str]) -> str:
+    def _get_endpoint(self, endpoint: UndefinedOr[str] | None) -> str | None:
         """Retrieves the API endpoint.
 
         If the endpoint is defined, it will be returned. Otherwise, it checks for

src/yandex_cloud_ml_sdk/exceptions.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
 from ._exceptions import (
-    AioRpcError, AsyncOperationError, DatasetValidationError, RunError, TuningError, WrongAsyncOperationStatusError,
-    YCloudMLError
+    AioRpcError, AsyncOperationError, DatasetValidationError, RunError, TuningError, UnknownEndpointError,
+    WrongAsyncOperationStatusError, YCloudMLError
 )
 
 __all__ = [
@@ -13,4 +13,5 @@
     'TuningError',
     'WrongAsyncOperationStatusError',
     'YCloudMLError',
+    'UnknownEndpointError',
 ]