Use Lease related methods from client service

Author: NickCao

packages/jumpstarter/conftest.py

@@ -16,7 +16,6 @@
 from jumpstarter_protocol import (
     jumpstarter_pb2,
     jumpstarter_pb2_grpc,
-    kubernetes_pb2,
     router_pb2_grpc,
 )
 
@@ -59,60 +58,6 @@ async def Register(self, request, context):
     async def Unregister(self, request, context):
         return jumpstarter_pb2.UnregisterResponse()
 
-    async def RequestLease(self, request, context):
-        lease_name = str(uuid4())
-        if "unsatisfiable" in request.selector.match_labels:
-            self.leases[lease_name] = "unsatisfiable"
-        else:
-            self.leases[lease_name] = 0
-        await self.status[0].send(
-            jumpstarter_pb2.StatusResponse(leased=True, lease_name=lease_name, client_name="dummy")
-        )
-        return jumpstarter_pb2.RequestLeaseResponse(name=lease_name)
-
-    async def GetLease(self, request, context):
-        if self.leases[request.name] == "unsatisfiable":
-            return jumpstarter_pb2.GetLeaseResponse(
-                exporter_uuid=str(uuid4()),
-                conditions=[
-                    kubernetes_pb2.Condition(
-                        type="Unsatisfiable",
-                        status="True",
-                    ),
-                ],
-            )
-
-        self.leases[request.name] += 1
-        if self.leases[request.name] == 1:
-            return jumpstarter_pb2.GetLeaseResponse(
-                exporter_uuid=str(uuid4()),
-                conditions=[],
-            )
-        if self.leases[request.name] == 2:
-            return jumpstarter_pb2.GetLeaseResponse(
-                exporter_uuid=str(uuid4()),
-                conditions=[
-                    kubernetes_pb2.Condition(
-                        type="Pending",
-                        status="True",
-                    ),
-                ],
-            )
-        else:
-            return jumpstarter_pb2.GetLeaseResponse(
-                exporter_uuid=str(uuid4()),
-                conditions=[
-                    kubernetes_pb2.Condition(
-                        type="Ready",
-                        status="True",
-                    ),
-                ],
-            )
-
-    async def ReleaseLease(self, request, context):
-        await self.status[0].send(jumpstarter_pb2.StatusResponse(leased=False))
-        return jumpstarter_pb2.ReleaseLeaseResponse()
-
     async def Dial(self, request, context):
         token = str(uuid4())
         await self.queue[0].send(token)

packages/jumpstarter/jumpstarter/client/lease.py

@@ -7,15 +7,16 @@
     contextmanager,
 )
 from dataclasses import dataclass, field
+from datetime import timedelta
 
 from anyio import fail_after, sleep
 from anyio.from_thread import BlockingPortal
-from google.protobuf import duration_pb2
 from grpc.aio import Channel
-from jumpstarter_protocol import jumpstarter_pb2, jumpstarter_pb2_grpc, kubernetes_pb2
+from jumpstarter_protocol import jumpstarter_pb2, jumpstarter_pb2_grpc
 
 from .exceptions import LeaseError
 from jumpstarter.client import client_from_path
+from jumpstarter.client.grpc import ClientService
 from jumpstarter.common import MetadataFilter, TemporaryUnixListener
 from jumpstarter.common.condition import condition_false, condition_message, condition_present_and_equal, condition_true
 from jumpstarter.common.grpc import translate_grpc_exceptions
@@ -31,6 +32,7 @@ class Lease(AbstractContextManager, AbstractAsyncContextManager):
     timeout: int = 1800
     metadata_filter: MetadataFilter = field(default_factory=MetadataFilter)
     portal: BlockingPortal
+    namespace: str
     name: str | None = field(default=None)
     allow: list[str]
     unsafe: bool
@@ -43,24 +45,23 @@ def __post_init__(self):
             super().__post_init__()
 
         self.controller = jumpstarter_pb2_grpc.ControllerServiceStub(self.channel)
+        self.svc = ClientService(channel=self.channel)
         self.manager = self.portal.wrap_async_context_manager(self)
 
     async def _create(self):
-        duration = duration_pb2.Duration()
-        duration.FromSeconds(self.timeout)
-        duration_str = str(duration.ToTimedelta())
+        duration = timedelta(seconds=self.timeout)
+        selector = ",".join(("{}={}".format(label[0], label[1]) for label in self.metadata_filter.labels.items()))
 
-        logger.debug("Creating lease request for labels %s for %s", self.metadata_filter.labels, duration_str)
+        logger.debug("Creating lease request for selector %s for duration %s", selector, duration)
         with translate_grpc_exceptions():
             self.name = (
-                await self.controller.RequestLease(
-                    jumpstarter_pb2.RequestLeaseRequest(
-                        duration=duration,
-                        selector=kubernetes_pb2.LabelSelector(match_labels=self.metadata_filter.labels),
-                    )
+                await self.svc.CreateLease(
+                    namespace=self.namespace,
+                    selector=selector,
+                    duration=timedelta(seconds=self.timeout),
                 )
             ).name
-        logger.info("Created lease request for labels %s for %s", self.metadata_filter.labels, duration_str)
+        logger.info("Created lease request for selector %s for duration %s", selector, duration)
 
     def request(self):
         """Request a lease, or verifies a lease which was already created.
@@ -97,7 +98,10 @@ async def _acquire(self):
             while True:
                 logger.debug("Polling Lease %s", self.name)
                 with translate_grpc_exceptions():
-                    result = await self.controller.GetLease(jumpstarter_pb2.GetLeaseRequest(name=self.name))
+                    result = await self.svc.GetLease(
+                        namespace=self.namespace,
+                        name=self.name,
+                    )
 
                 # lease ready
                 if condition_true(result.conditions, "Ready"):
@@ -131,7 +135,10 @@ async def __aenter__(self):
     async def __aexit__(self, exc_type, exc_value, traceback):
         if self.release:
             logger.info("Releasing Lease %s", self.name)
-            await self.controller.ReleaseLease(jumpstarter_pb2.ReleaseLeaseRequest(name=self.name))
+            await self.svc.DeleteLease(
+                namespace=self.namespace,
+                name=self.name,
+            )
 
     def __enter__(self):
         # wraps the async context manager enter

packages/jumpstarter/jumpstarter/config/client.py

@@ -6,7 +6,6 @@
 import grpc
 import yaml
 from anyio.from_thread import BlockingPortal, start_blocking_portal
-from jumpstarter_protocol import jumpstarter_pb2, jumpstarter_pb2_grpc
 from pydantic import BaseModel, ConfigDict, Field, ValidationError
 
 from .common import CONFIG_PATH, ObjectMeta
@@ -118,6 +117,7 @@ async def request_lease_async(
 
         lease = Lease(
             channel=await self.channel(),
+            namespace=self.metadata.namespace,
             name=None,
             metadata_filter=metadata_filter,
             portal=portal,
@@ -129,14 +129,15 @@ async def request_lease_async(
             return await lease.request_async()
 
     async def list_leases_async(self):
-        controller = jumpstarter_pb2_grpc.ControllerServiceStub(await self.channel())
+        svc = ClientService(channel=await self.channel())
         with translate_grpc_exceptions():
-            return (await controller.ListLeases(jumpstarter_pb2.ListLeasesRequest())).names
+            result = await svc.ListLeases(namespace=self.metadata.namespace)
+            return [lease.name for lease in result.leases]
 
     async def release_lease_async(self, name):
-        controller = jumpstarter_pb2_grpc.ControllerServiceStub(await self.channel())
+        svc = ClientService(channel=await self.channel())
         with translate_grpc_exceptions():
-            await controller.ReleaseLease(jumpstarter_pb2.ReleaseLeaseRequest(name=name))
+            await svc.DeleteLease(namespace=self.metadata.namespace, name=name)
 
     @asynccontextmanager
     async def lease_async(
@@ -154,6 +155,7 @@ async def lease_async(
 
         async with Lease(
             channel=await self.channel(),
+            namespace=self.metadata.namespace,
             name=lease_name,
             metadata_filter=metadata_filter,
             portal=portal,

packages/jumpstarter/jumpstarter/config/exporter_test.py

@@ -1,70 +1,14 @@
-from contextlib import ExitStack
 from pathlib import Path
 
 import pytest
-from anyio import create_task_group
-from anyio.from_thread import start_blocking_portal
 
-from .client import ClientConfigV1Alpha1, ClientConfigV1Alpha1Drivers
 from .common import ObjectMeta
 from .exporter import ExporterConfigV1Alpha1, ExporterConfigV1Alpha1DriverInstance
 from .tls import TLSConfigV1Alpha1
-from jumpstarter.common import MetadataFilter
 
 pytestmark = pytest.mark.anyio
 
 
-async def test_exporter_serve(mock_controller):
-    exporter = ExporterConfigV1Alpha1(
-        apiVersion="jumpstarter.dev/v1alpha1",
-        kind="ExporterConfig",
-        metadata=ObjectMeta(namespace="default", name="test"),
-        endpoint=mock_controller,
-        token="dummy-exporter-token",
-        export={
-            "power": ExporterConfigV1Alpha1DriverInstance(
-                type="jumpstarter_driver_power.driver.MockPower",
-            ),
-            "nested": ExporterConfigV1Alpha1DriverInstance(
-                children={
-                    "tcp": ExporterConfigV1Alpha1DriverInstance(
-                        type="jumpstarter_driver_network.driver.TcpNetwork",
-                        config={
-                            "host": "127.0.0.1",
-                            "port": 8080,
-                        },
-                    )
-                }
-            ),
-        },
-    )
-
-    client = ClientConfigV1Alpha1(
-        name="testclient",
-        metadata=ObjectMeta(namespace="default", name="testclient"),
-        endpoint=mock_controller,
-        token="dummy-client-token",
-        drivers=ClientConfigV1Alpha1Drivers(allow=[], unsafe=True),
-        tls=TLSConfigV1Alpha1(insecure=True),
-    )
-
-    async with create_task_group() as tg:
-        tg.start_soon(exporter.serve)
-
-        with start_blocking_portal() as portal:
-            async with client.lease_async(
-                metadata_filter=MetadataFilter(),
-                lease_name=None,
-                portal=portal,
-            ) as lease:
-                with ExitStack() as stack:
-                    async with lease.connect_async(stack) as client:
-                        await client.power.call_async("on")
-                        assert hasattr(client.nested, "tcp")
-
-        tg.cancel_scope.cancel()
-
-
 def test_exporter_config(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
     monkeypatch.setattr(ExporterConfigV1Alpha1, "BASE_PATH", tmp_path)