Use smart stub in clients

Author: NickCao

packages/jumpstarter/jumpstarter/client/client.py

@@ -6,9 +6,8 @@
 import grpc
 from anyio.from_thread import BlockingPortal
 from google.protobuf import empty_pb2
-from jumpstarter_protocol import jumpstarter_pb2_grpc
 
-from .grpc import SmartExporterServiceStub
+from .grpc import SmartExporterStub
 from jumpstarter.client import DriverClient
 from jumpstarter.common.importlib import import_class
 
@@ -34,30 +33,24 @@ async def client_from_channel(
     reports = {}
     clients = OrderedDict()
 
-    response = await SmartExporterServiceStub([channel]).GetReport(empty_pb2.Empty())
+    response = await SmartExporterStub([channel]).GetReport(empty_pb2.Empty())
 
+    channels = [channel]
     if use_alternative_endpoints:
         for endpoint in response.alternative_endpoints:
             if endpoint.certificate:
-                attempted_channel = grpc.aio.secure_channel(
-                    endpoint.endpoint,
-                    grpc.ssl_channel_credentials(
-                        root_certificates=endpoint.certificate.encode(),
-                        private_key=endpoint.client_private_key.encode(),
-                        certificate_chain=endpoint.client_certificate.encode(),
-                    ),
-                )
-                try:
-                    response = await jumpstarter_pb2_grpc.ExporterServiceStub(attempted_channel).GetReport(
-                        empty_pb2.Empty()
+                channels.append(
+                    grpc.aio.secure_channel(
+                        endpoint.endpoint,
+                        grpc.ssl_channel_credentials(
+                            root_certificates=endpoint.certificate.encode(),
+                            private_key=endpoint.client_private_key.encode(),
+                            certificate_chain=endpoint.client_certificate.encode(),
+                        ),
                     )
-                except Exception:
-                    pass  # TODO: log failed attempt
-                else:
-                    channel = attempted_channel
-                    break
-            else:
-                continue
+                )
+
+    stub = SmartExporterStub(list(reversed(channels)))
 
     for index, report in enumerate(response.reports):
         topo[index] = []
@@ -77,7 +70,7 @@ async def client_from_channel(
         client = client_class(
             uuid=UUID(report.uuid),
             labels=report.labels,
-            channel=channel,
+            stub=stub,
             portal=portal,
             stack=stack.enter_context(ExitStack()),
             children={reports[k].labels["jumpstarter.dev/name"]: clients[k] for k in topo[index]},

packages/jumpstarter/jumpstarter/client/core.py

@@ -5,11 +5,12 @@
 import logging
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
+from typing import Any
 
 from anyio import create_task_group
 from google.protobuf import empty_pb2
 from grpc import StatusCode
-from grpc.aio import AioRpcError, Channel
+from grpc.aio import AioRpcError
 from jumpstarter_protocol import jumpstarter_pb2, jumpstarter_pb2_grpc, router_pb2_grpc
 
 from jumpstarter.common import Metadata
@@ -60,16 +61,14 @@ class AsyncDriverClient(
     Backing implementation of blocking driver client.
     """
 
-    channel: Channel
+    stub: Any
 
     log_level: str = "INFO"
     logger: logging.Logger = field(init=False)
 
     def __post_init__(self):
         if hasattr(super(), "__post_init__"):
             super().__post_init__()
-        jumpstarter_pb2_grpc.ExporterServiceStub.__init__(self, self.channel)
-        router_pb2_grpc.RouterServiceStub.__init__(self, self.channel)
         self.logger = logging.getLogger(self.__class__.__name__)
         self.logger.setLevel(self.log_level)
 
@@ -89,7 +88,7 @@ async def call_async(self, method, *args):
         )
 
         try:
-            response = await self.DriverCall(request)
+            response = await self.stub.DriverCall(request)
         except AioRpcError as e:
             match e.code():
                 case StatusCode.UNIMPLEMENTED:
@@ -113,7 +112,7 @@ async def streamingcall_async(self, method, *args):
         )
 
         try:
-            async for response in self.StreamingDriverCall(request):
+            async for response in self.stub.StreamingDriverCall(request):
                 yield decode_value(response.result)
         except AioRpcError as e:
             match e.code():
@@ -128,7 +127,7 @@ async def streamingcall_async(self, method, *args):
 
     @asynccontextmanager
     async def stream_async(self, method):
-        context = self.Stream(
+        context = self.stub.Stream(
             metadata=StreamRequestMetadata.model_construct(request=DriverStreamRequest(uuid=self.uuid, method=method))
             .model_dump(mode="json", round_trip=True)
             .items(),
@@ -142,7 +141,7 @@ async def resource_async(
         self,
         stream,
     ):
-        context = self.Stream(
+        context = self.stub.Stream(
             metadata=StreamRequestMetadata.model_construct(request=ResourceStreamRequest(uuid=self.uuid))
             .model_dump(mode="json", round_trip=True)
             .items(),
@@ -160,7 +159,7 @@ def __log(self, level: int, msg: str):
     @asynccontextmanager
     async def log_stream_async(self):
         async def log_stream():
-            async for response in self.LogStream(empty_pb2.Empty()):
+            async for response in self.stub.LogStream(empty_pb2.Empty()):
                 self.__log(logging.getLevelName(response.severity), response.message)
 
         async with create_task_group() as tg:

packages/jumpstarter/jumpstarter/client/grpc.py

@@ -3,8 +3,8 @@
 from collections import OrderedDict
 from dataclasses import InitVar, dataclass, field
 from datetime import datetime, timedelta
-from functools import partial
-from typing import Generic, Type, TypeVar
+from types import SimpleNamespace
+from typing import Any
 
 import yaml
 from google.protobuf import duration_pb2, field_mask_pb2, json_format
@@ -256,19 +256,18 @@ async def DeleteLease(self, *, name: str):
             )
 
 
-T = TypeVar("T")
-
-
 @dataclass(frozen=True, slots=True)
-class SmartStub(Generic[T]):
+class SmartExporterStub:
     channels: InitVar[list[Channel]]
-    cls: InitVar[Type]
 
-    __stubs: dict[Channel, T] = field(init=False, default_factory=OrderedDict)
+    __stubs: dict[Channel, Any] = field(init=False, default_factory=OrderedDict)
 
-    def __post_init__(self, channels, cls):
+    def __post_init__(self, channels):
         for channel in channels:
-            self.__stubs[channel] = cls(channel)
+            stub = SimpleNamespace()
+            jumpstarter_pb2_grpc.ExporterServiceStub.__init__(stub, channel)
+            router_pb2_grpc.RouterServiceStub.__init__(stub, channel)
+            self.__stubs[channel] = stub
 
     def __getattr__(self, name):
         for channel, stub in self.__stubs.items():
@@ -277,13 +276,3 @@ def __getattr__(self, name):
                 return getattr(stub, name)
         # or fallback to the last channel (via router)
         return getattr(next(reversed(self.__stubs.values())), name)
-
-
-SmartExporterServiceStub = partial(
-    SmartStub[jumpstarter_pb2_grpc.ExporterServiceStub],
-    cls=jumpstarter_pb2_grpc.ExporterServiceStub,
-)
-SmartRouterServiceStub = partial(
-    SmartStub[router_pb2_grpc.RouterServiceStub],
-    cls=router_pb2_grpc.RouterServiceStub,
-)