Merge pull request #36 from jumpstarter-dev/opendal

Implement file access with opendal

Author: mangelajo

jumpstarter/common/aiohttp.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+
+from aiohttp import StreamReader
+from anyio import EndOfStream
+from anyio.abc import ByteReceiveStream
+
+
+@dataclass(kw_only=True)
+class AiohttpStream(ByteReceiveStream):
+    stream: StreamReader
+
+    async def receive(self, max_bytes=65536):
+        data = await self.stream.read(n=max_bytes)
+        if len(data) == 0:
+            raise EndOfStream
+        return data
+
+    async def aclose(self):
+        pass

jumpstarter/common/opendal.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+
+from anyio import EndOfStream
+from anyio.abc import ByteReceiveStream
+from opendal import AsyncFile
+
+
+@dataclass(kw_only=True)
+class AsyncFileStream(ByteReceiveStream):
+    file: AsyncFile
+
+    async def receive(self, max_bytes=65536):
+        data = await self.file.read(size=max_bytes)
+        if len(data) == 0:
+            raise EndOfStream
+        return data
+
+    async def aclose(self):
+        await self.file.close()

jumpstarter/drivers/base.py

@@ -3,14 +3,17 @@
 """
 
 from abc import ABCMeta, abstractmethod
+from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from typing import Any
 from uuid import UUID, uuid4
 
+import aiohttp
 from anyio.from_thread import BlockingPortal
 from grpc import StatusCode
 
 from jumpstarter.common import Metadata
+from jumpstarter.common.aiohttp import AiohttpStream
 from jumpstarter.common.streams import (
     create_memory_stream,
     forward_server_stream,
@@ -22,6 +25,8 @@
     MARKER_STREAMCALL,
     MARKER_STREAMING_DRIVERCALL,
 )
+from jumpstarter.drivers.resources import ClientStreamResource, PresignedRequestResource, Resource
+from jumpstarter.drivers.streams import DriverStreamRequest, ResourceStreamRequest, StreamRequest
 from jumpstarter.v1 import jumpstarter_pb2, jumpstarter_pb2_grpc, router_pb2_grpc
 
 
@@ -83,14 +88,16 @@ async def Stream(self, request_iterator, context):
         """
         metadata = dict(context.invocation_metadata())
 
-        match metadata["kind"]:
-            case "connect":
-                method = await self.__lookup_drivercall(metadata["method"], context, MARKER_STREAMCALL)
+        request = StreamRequest.validate_json(metadata["request"], strict=True)
+
+        match request:
+            case DriverStreamRequest(method=driver_method):
+                method = await self.__lookup_drivercall(driver_method, context, MARKER_STREAMCALL)
 
                 async for v in method(request_iterator, context):
                     yield v
 
-            case "resource":
+            case ResourceStreamRequest():
                 remote, resource = create_memory_stream()
 
                 resource_uuid = uuid4()
@@ -133,6 +140,16 @@ def items(self, parent=None):
 
         return [(self.uuid, parent.uuid if parent else None, self)]
 
+    @asynccontextmanager
+    async def resource(self, handle: str):
+        handle = Resource.validate_python(handle)
+        match handle:
+            case ClientStreamResource(uuid=uuid):
+                yield self.resources[uuid]
+            case PresignedRequestResource(headers=headers, url=url, method=method):
+                async with aiohttp.request(method, url, headers=headers, raise_for_status=True) as resp:
+                    yield AiohttpStream(stream=resp.content)
+
     async def __lookup_drivercall(self, name, context, marker):
         """Lookup drivercall by method name
 

jumpstarter/drivers/core.py

@@ -4,18 +4,23 @@
 
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
+from uuid import UUID
 
 from anyio import create_task_group, sleep_forever
 from anyio.streams.stapled import StapledObjectStream
 from google.protobuf import json_format, struct_pb2
 from grpc.aio import Channel
+from opendal import AsyncOperator
 
 from jumpstarter.common import Metadata
+from jumpstarter.common.opendal import AsyncFileStream
 from jumpstarter.common.progress import ProgressStream
 from jumpstarter.common.streams import (
     create_memory_stream,
     forward_client_stream,
 )
+from jumpstarter.drivers.resources import ClientStreamResource, PresignedRequestResource
+from jumpstarter.drivers.streams import DriverStreamRequest, ResourceStreamRequest
 from jumpstarter.v1 import jumpstarter_pb2, jumpstarter_pb2_grpc, router_pb2_grpc
 
 
@@ -68,7 +73,7 @@ async def stream_async(self, method):
         async with forward_client_stream(
             self,
             device_stream,
-            {"kind": "connect", "uuid": str(self.uuid), "method": method}.items(),
+            {"request": DriverStreamRequest(uuid=self.uuid, method=method).model_dump_json()}.items(),
         ):
             async with client_stream:
                 yield client_stream
@@ -80,7 +85,7 @@ async def handle(client):
                 async with forward_client_stream(
                     self,
                     client,
-                    {"kind": "connect", "uuid": str(self.uuid), "method": method}.items(),
+                    {"request": DriverStreamRequest(uuid=self.uuid, method=method).model_dump_json()}.items(),
                 ):
                     await sleep_forever()
 
@@ -104,6 +109,22 @@ async def resource_async(
             async with forward_client_stream(
                 self,
                 combined,
-                {"kind": "resource", "uuid": str(self.uuid)}.items(),
+                {"request": ResourceStreamRequest(uuid=self.uuid).model_dump_json()}.items(),
             ):
-                yield (await rx.receive()).decode()
+                yield ClientStreamResource(uuid=UUID((await rx.receive()).decode())).model_dump(mode="json")
+
+    @asynccontextmanager
+    async def file_async(
+        self,
+        operator: AsyncOperator,
+        path: str,
+    ):
+        if operator.capability().presign:
+            presigned = await operator.presign_read(path, expire_second=60)
+            yield PresignedRequestResource(
+                headers=presigned.headers, url=presigned.url, method=presigned.method
+            ).model_dump(mode="json")
+        else:
+            file = await operator.open(path, "rb")
+            async with self.resource_async(AsyncFileStream(file=file)) as handle:
+                yield handle

jumpstarter/drivers/dutlink/base.py

@@ -2,7 +2,6 @@
 from collections.abc import AsyncGenerator
 from dataclasses import dataclass, field
 from pathlib import Path
-from uuid import UUID
 
 import pyudev
 import usb.core
@@ -103,8 +102,9 @@ async def write(self, src: str):
                 await sleep(1)
 
         async with await FileWriteStream.from_path(self.storage_device) as stream:
-            async for chunk in self.resources[UUID(src)]:
-                await stream.send(chunk)
+            async with self.resource(src) as res:
+                async for chunk in res:
+                    await stream.send(chunk)
 
 
 @dataclass(kw_only=True)

jumpstarter/drivers/mixins.py

@@ -11,7 +11,7 @@
 
 from anyio import create_unix_listener
 from anyio.from_thread import BlockingPortal
-from anyio.streams.file import FileReadStream
+from opendal import Operator
 from pexpect.fdpexpect import fdspawn
 
 
@@ -70,15 +70,6 @@ class ResourceMixin:
     """Resource"""
 
     @contextmanager
-    def local_file(
-        self,
-        filepath,
-    ):
-        """
-        Share local file with driver
-
-        :param str filepath: path to file
-        """
-        with self.portal.wrap_async_context_manager(self.portal.call(FileReadStream.from_path, filepath)) as file:
-            with self.portal.wrap_async_context_manager(self.resource_async(file)) as uuid:
-                yield uuid
+    def file(self, operator: Operator, path: str):
+        with self.portal.wrap_async_context_manager(self.file_async(operator.to_async_operator(), path)) as uuid:
+            yield uuid

jumpstarter/drivers/resources.py

@@ -0,0 +1,24 @@
+from typing import Annotated, Literal, Union
+from uuid import UUID
+
+from pydantic import BaseModel, Field, TypeAdapter
+
+
+class ClientStreamResource(BaseModel):
+    kind: Literal["client_stream"] = "client_stream"
+    uuid: UUID
+
+
+class PresignedRequestResource(BaseModel):
+    kind: Literal["presigned_request"] = "presigned_request"
+    headers: dict[str, str]
+    url: str
+    method: str
+
+
+Resource = TypeAdapter(
+    Annotated[
+        Union[ClientStreamResource, PresignedRequestResource],
+        Field(discriminator="kind"),
+    ]
+)

jumpstarter/drivers/storage/base.py

@@ -1,9 +1,9 @@
 from abc import ABCMeta, abstractmethod
 from tempfile import NamedTemporaryFile
-from uuid import UUID
 
 import click
 from anyio.streams.file import FileWriteStream
+from opendal import Operator
 
 from jumpstarter.drivers import Driver, DriverClient, export
 from jumpstarter.drivers.mixins import ResourceMixin
@@ -40,8 +40,12 @@ def off(self):
     def write(self, handle):
         return self.call("write", handle)
 
+    def write_file(self, operator: Operator, path: str):
+        with self.file(operator, path) as handle:
+            return self.call("write", handle)
+
     def write_local_file(self, filepath):
-        with self.local_file(filepath) as handle:
+        with self.file(Operator("fs", root="/"), filepath) as handle:
             return self.call("write", handle)
 
     def cli(self):
@@ -90,5 +94,6 @@ async def off(self):
     async def write(self, src: str):
         with NamedTemporaryFile() as file:
             async with FileWriteStream(file) as stream:
-                async for chunk in self.resources[UUID(src)]:
-                    await stream.send(chunk)
+                async with self.resource(src) as res:
+                    async for chunk in res:
+                        await stream.send(chunk)

jumpstarter/drivers/storage/base_test.py

@@ -1,19 +1,36 @@
-import os
-from tempfile import NamedTemporaryFile
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import pytest
+from opendal import Operator
 
 from jumpstarter.common.utils import serve
 from jumpstarter.drivers.storage import MockStorageMux
 
 
-def test_drivers_mock_storage_mux():
+def test_drivers_mock_storage_mux_fs():
+    with serve(MockStorageMux(name="storage")) as client:
+        with TemporaryDirectory() as tempdir:
+            fs = Operator("fs", root=tempdir)
+
+            fs.write("test", b"testcontent" * 1000)
+
+            client.write_file(fs, "test")
+            client.write_local_file(str(Path(tempdir) / "test"))
+
+
+@pytest.mark.skip(reason="require minio")
+def test_drivers_mock_storage_mux_s3():
     with serve(MockStorageMux(name="storage")) as client:
-        with NamedTemporaryFile(delete=False) as file:
-            file.write(b"testcontent" * 1000)
-            file.close()
+        s3 = Operator(
+            "s3",
+            bucket="test",
+            endpoint="http://127.0.0.1:9000",
+            region="us-east-1",
+            access_key_id="minioadmin",
+            secret_access_key="minioadmin",
+        )
 
-            client.off()
-            client.dut()
-            client.host()
-            client.write_local_file(file.name)
+        s3.write("test", b"testcontent" * 1000)
 
-            os.unlink(file.name)
+        client.write_file(s3, "test")

jumpstarter/drivers/streams.py

@@ -0,0 +1,23 @@
+from typing import Annotated, Literal, Union
+from uuid import UUID
+
+from pydantic import BaseModel, Field, TypeAdapter
+
+
+class ResourceStreamRequest(BaseModel):
+    kind: Literal["resource"] = "resource"
+    uuid: UUID
+
+
+class DriverStreamRequest(BaseModel):
+    kind: Literal["driver"] = "driver"
+    uuid: UUID
+    method: str
+
+
+StreamRequest = TypeAdapter(
+    Annotated[
+        Union[ResourceStreamRequest, DriverStreamRequest],
+        Field(discriminator="kind"),
+    ]
+)

jumpstarter/exporter/session.py

@@ -3,6 +3,7 @@
 
 from jumpstarter.common import Metadata
 from jumpstarter.drivers.base import Driver
+from jumpstarter.drivers.streams import StreamRequest
 from jumpstarter.v1 import (
     jumpstarter_pb2,
     jumpstarter_pb2_grpc,
@@ -51,5 +52,7 @@ async def StreamingDriverCall(self, request, context):
     async def Stream(self, request_iterator, context):
         metadata = dict(context.invocation_metadata())
 
-        async for v in self[UUID(metadata["uuid"])].Stream(request_iterator, context):
+        request = StreamRequest.validate_json(metadata["request"], strict=True)
+
+        async for v in self[request.uuid].Stream(request_iterator, context):
             yield v