Merge branch 'main' into client-svc-lease

Author: NickCao

docs/source/cli/clients.md

@@ -32,6 +32,9 @@ metadata:
   name: john
 endpoint: grpc.jumpstarter.192.168.1.10.nip.io:8082
 token: <<token>>
+grpcConfig:
+  # please refer to the https://grpc.github.io/grpc/core/group__grpc__arg__keys.html documentation
+  grpc.keepalive_time_ms: 20000
 tls:
   ca: ''
   insecure: False

docs/source/introduction/exporters.md

@@ -30,6 +30,9 @@ metadata:
   name: demo
 endpoint: grpc.jumpstarter.example.com:443
 token: xxxxx
+grpcConfig:
+    # Please refer to the https://grpc.github.io/grpc/core/group__grpc__arg__keys.html documentation
+    grpc.keepalive_time_ms: 20000
 export:
   power:
     type: jumpstarter_driver_yepkit.driver.Ykush

packages/jumpstarter-cli-admin/jumpstarter_cli_admin/get_test.py

@@ -765,6 +765,7 @@ async def test_get_exporters(_load_kube_config_mock, list_exporters_mock: AsyncM
 
 EXPORTERS_DEVICES_LIST_NAME = EXPORTERS_LIST_NAME
 
+
 @pytest.mark.anyio
 @patch.object(ExportersV1Alpha1Api, "list_exporters")
 @patch.object(ExportersV1Alpha1Api, "_load_kube_config")
@@ -1151,6 +1152,7 @@ async def test_get_lease(_load_kube_config_mock, get_lease_mock: AsyncMock):
 lease.jumpstarter.dev/82a8ac0d-d7ff-4009-8948-18a3c5c607b2
 """
 
+
 @pytest.mark.anyio
 @patch.object(LeasesV1Alpha1Api, "list_leases")
 @patch.object(LeasesV1Alpha1Api, "_load_kube_config")

packages/jumpstarter-cli-client/jumpstarter_cli_client/client_shell.py

@@ -22,6 +22,7 @@ def client_shell(config, selector: str, lease_name):
         metadata_filter=MetadataFilter(labels=selector_to_labels(selector)), lease_name=lease_name
     ) as lease:
         with lease.serve_unix() as path:
-            exit_code = launch_shell(path, "remote", config.drivers.allow, config.drivers.unsafe)
+            with lease.monitor():
+                exit_code = launch_shell(path, "remote", config.drivers.allow, config.drivers.unsafe)
 
     sys.exit(exit_code)

packages/jumpstarter/jumpstarter/client/grpc.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from datetime import timedelta
+from datetime import datetime, timedelta
 
 import yaml
 from google.protobuf import duration_pb2, field_mask_pb2, json_format
@@ -52,6 +52,7 @@ class Lease(BaseModel):
     client: str
     exporter: str
     conditions: list[kubernetes_pb2.Condition]
+    effective_begin_time: datetime | None = None
 
     model_config = ConfigDict(
         arbitrary_types_allowed=True,
@@ -72,13 +73,20 @@ def from_protobuf(cls, data: client_pb2.Lease) -> Lease:
         else:
             exporter = ""
 
+        effective_begin_time = None
+        if data.effective_begin_time:
+            effective_begin_time = data.effective_begin_time.ToDatetime(
+                tzinfo=datetime.now().astimezone().tzinfo,
+            )
+
         return cls(
             namespace=namespace,
             name=name,
             selector=data.selector,
             duration=data.duration.ToTimedelta(),
             client=client,
             exporter=exporter,
+            effective_begin_time=effective_begin_time,
             conditions=data.conditions,
         )
 

packages/jumpstarter/jumpstarter/client/lease.py

@@ -7,9 +7,10 @@
     contextmanager,
 )
 from dataclasses import dataclass, field
-from datetime import timedelta
+from datetime import datetime, timedelta
+from typing import Any
 
-from anyio import fail_after, sleep
+from anyio import create_task_group, fail_after, sleep
 from anyio.from_thread import BlockingPortal
 from grpc.aio import Channel
 from jumpstarter_protocol import jumpstarter_pb2, jumpstarter_pb2_grpc
@@ -39,6 +40,7 @@ class Lease(AbstractContextManager, AbstractAsyncContextManager):
     release: bool = True  # release on contexts exit
     controller: jumpstarter_pb2_grpc.ControllerServiceStub = field(init=False)
     tls_config: TLSConfigV1Alpha1 = field(default_factory=TLSConfigV1Alpha1)
+    grpc_options: dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         if hasattr(super(), "__post_init__"):
@@ -62,6 +64,11 @@ async def _create(self):
             ).name
         logger.info("Created lease request for selector %s for duration %s", selector, duration)
 
+    async def get(self):
+        with translate_grpc_exceptions():
+            svc = ClientService(channel=self.channel, namespace=self.namespace)
+            return await svc.GetLease(name=self.name)
+
     def request(self):
         """Request a lease, or verifies a lease which was already created.
 
@@ -96,11 +103,7 @@ async def _acquire(self):
         with fail_after(300):  # TODO: configurable timeout
             while True:
                 logger.debug("Polling Lease %s", self.name)
-                with translate_grpc_exceptions():
-                    result = await self.svc.GetLease(
-                        name=self.name,
-                    )
-
+                result = await self.get()
                 # lease ready
                 if condition_true(result.conditions, "Ready"):
                     logger.debug("Lease %s acquired", self.name)
@@ -148,14 +151,39 @@ def __exit__(self, exc_type, exc_value, traceback):
     async def handle_async(self, stream):
         logger.debug("Connecting to Lease with name %s", self.name)
         response = await self.controller.Dial(jumpstarter_pb2.DialRequest(lease_name=self.name))
-        async with connect_router_stream(response.router_endpoint, response.router_token, stream, self.tls_config):
+        async with connect_router_stream(
+            response.router_endpoint, response.router_token, stream, self.tls_config, self.grpc_options
+        ):
             pass
 
     @asynccontextmanager
     async def serve_unix_async(self):
         async with TemporaryUnixListener(self.handle_async) as path:
             yield path
 
+    @asynccontextmanager
+    async def monitor_async(self, threshold: timedelta = timedelta(minutes=5)):
+        async def _monitor():
+            while True:
+                lease = await self.get()
+                if lease.effective_begin_time:
+                    end_time = lease.effective_begin_time + lease.duration
+                    remain = end_time - datetime.now(tz=datetime.now().astimezone().tzinfo)
+                    if remain < threshold:
+                        logger.info("Lease {} ending soon in {} at {}".format(self.name, remain, end_time))
+                        await sleep(threshold.total_seconds())
+                    else:
+                        await sleep(5)
+                else:
+                    await sleep(1)
+
+        async with create_task_group() as tg:
+            tg.start_soon(_monitor)
+            try:
+                yield
+            finally:
+                tg.cancel_scope.cancel()
+
     @asynccontextmanager
     async def connect_async(self, stack):
         async with self.serve_unix_async() as path:
@@ -172,3 +200,8 @@ def connect(self):
     def serve_unix(self):
         with self.portal.wrap_async_context_manager(self.serve_unix_async()) as path:
             yield path
+
+    @contextmanager
+    def monitor(self, threshold: timedelta = timedelta(minutes=5)):
+        with self.portal.wrap_async_context_manager(self.monitor_async(threshold)):
+            yield

packages/jumpstarter/jumpstarter/common/grpc.py

@@ -3,6 +3,7 @@
 import socket
 import ssl
 from contextlib import contextmanager
+from typing import Any, Sequence, Tuple
 from urllib.parse import urlparse
 
 import grpc
@@ -34,20 +35,28 @@ def ssl_channel_credentials(target: str, tls_config):
         return grpc.ssl_channel_credentials()
 
 
-def aio_secure_channel(target: str, credentials: grpc.ChannelCredentials):
+def aio_secure_channel(target: str, credentials: grpc.ChannelCredentials, grpc_options: dict[str, Any] | None):
     return grpc.aio.secure_channel(
         target,
         credentials,
-        options=(
-            ("grpc.lb_policy_name", "round_robin"),
-            ("grpc.keepalive_time_ms", 350000),
-            ("grpc.keepalive_timeout_ms", 5000),
-            ("grpc.http2.max_pings_without_data", 5),
-            ("grpc.keepalive_permit_without_calls", 1),
-        ),
+        options=_override_default_grpc_options(grpc_options),
     )
 
 
+def _override_default_grpc_options(grpc_options: dict[str, str | int] | None) -> Sequence[Tuple[str, Any]]:
+    defaults = (
+        ("grpc.lb_policy_name", "round_robin"),
+        # we keep a low keepalive time to avoid idle timeouts on cloud load balancers
+        ("grpc.keepalive_time_ms", 20000),
+        ("grpc.keepalive_timeout_ms", 5000),
+        ("grpc.http2.max_pings_without_data", 5),
+        ("grpc.keepalive_permit_without_calls", 1),
+    )
+    options = dict(defaults)
+    options.update(grpc_options or {})
+    return tuple(options.items())
+
+
 def configure_grpc_env():
     # disable informative logs by default, i.e.:
     # WARNING: All log messages before absl::InitializeLog() is called are written to STDERR

packages/jumpstarter/jumpstarter/common/streams.py

@@ -32,13 +32,13 @@ class StreamRequestMetadata(BaseModel):
 
 
 @asynccontextmanager
-async def connect_router_stream(endpoint, token, stream, tls_config):
+async def connect_router_stream(endpoint, token, stream, tls_config, grpc_options):
     credentials = grpc.composite_channel_credentials(
         ssl_channel_credentials(endpoint, tls_config),
         grpc.access_token_call_credentials(token),
     )
 
-    async with aio_secure_channel(endpoint, credentials) as channel:
+    async with aio_secure_channel(endpoint, credentials, grpc_options) as channel:
         router = router_pb2_grpc.RouterServiceStub(channel)
         context = router.Stream(metadata=())
         async with RouterStream(context=context) as s:

packages/jumpstarter/jumpstarter/config/client.py

@@ -49,6 +49,7 @@ class ClientConfigV1Alpha1(BaseModel):
     endpoint: str
     tls: TLSConfigV1Alpha1 = Field(default_factory=TLSConfigV1Alpha1)
     token: str
+    grpcOptions: dict[str, str | int] | None = Field(default_factory=dict)
 
     drivers: ClientConfigV1Alpha1Drivers
 
@@ -58,7 +59,7 @@ async def channel(self):
             call_credentials("Client", self.metadata, self.token),
         )
 
-        return aio_secure_channel(self.endpoint, credentials)
+        return aio_secure_channel(self.endpoint, credentials, self.grpcOptions)
 
     @contextmanager
     def lease(self, metadata_filter: MetadataFilter, lease_name: str | None = None):
@@ -161,6 +162,7 @@ async def request_lease_async(
             allow=self.drivers.allow,
             unsafe=self.drivers.unsafe,
             tls_config=self.tls,
+            grpc_options=self.grpcOptions,
         )
         with translate_grpc_exceptions():
             return await lease.request_async()
@@ -204,6 +206,7 @@ async def lease_async(
             unsafe=self.drivers.unsafe,
             release=release_lease,
             tls_config=self.tls,
+            grpc_options=self.grpcOptions,
         ) as lease:
             yield lease
 

packages/jumpstarter/jumpstarter/config/client_config_test.py

@@ -206,6 +206,7 @@ def test_client_config_save(monkeypatch: pytest.MonkeyPatch):
   ca: ''
   insecure: false
 token: dGhpc2lzYXRva2VuLTEyMzQxMjM0MTIzNEyMzQtc2Rxd3Jxd2VycXdlcnF3ZXJxd2VyLTEyMzQxMjM0MTIz
+grpcOptions: {}
 drivers:
   allow:
   - jumpstarter.drivers.*
@@ -241,6 +242,7 @@ def test_client_config_save_explicit_path():
   ca: ''
   insecure: false
 token: dGhpc2lzYXRva2VuLTEyMzQxMjM0MTIzNEyMzQtc2Rxd3Jxd2VycXdlcnF3ZXJxd2VyLTEyMzQxMjM0MTIz
+grpcOptions: {}
 drivers:
   allow:
   - jumpstarter.drivers.*
@@ -274,6 +276,7 @@ def test_client_config_save_unsafe_drivers():
   ca: ''
   insecure: false
 token: dGhpc2lzYXRva2VuLTEyMzQxMjM0MTIzNEyMzQtc2Rxd3Jxd2VycXdlcnF3ZXJxd2VyLTEyMzQxMjM0MTIz
+grpcOptions: {}
 drivers:
   allow: []
   unsafe: true