ellf-broker-sdk 7.0.8__tar.gz

ellf_broker_sdk-7.0.8/MANIFEST.in

@@ -0,0 +1 @@
+include requirements.in

ellf_broker_sdk-7.0.8/PKG-INFO

@@ -0,0 +1,18 @@
+Metadata-Version: 2.4
+Name: ellf-broker-sdk
+Version: 7.0.8
+Summary: Ellf Broker SDK
+Home-page: https://github.com/explosion/ellf
+Author: ExplosionAI GmbH
+Author-email: contact@explosion.ai
+License: All rights reserved
+Requires-Dist: httpx
+Requires-Dist: pydantic-settings>=2.0.0
+Requires-Dist: pydantic>=2.0.0
+Requires-Dist: pyjwt
+Dynamic: author
+Dynamic: author-email
+Dynamic: home-page
+Dynamic: license
+Dynamic: requires-dist
+Dynamic: summary

ellf_broker_sdk-7.0.8/README.md

@@ -0,0 +1,9 @@
+<a href="https://explosion.ai"><img src="https://explosion.ai/assets/img/logo.svg" width="125" height="125" align="right" /></a>
+
+# Ellf Broker SDK
+
+The `models.py` is auto-generated using the OpenAPI spec and [`datamodel-code-generator`](https://koxudaxi.github.io/datamodel-code-generator/). To generate the SDK, make sure you have all requirements installed and run the following from the project root directory:
+
+```bash
+pdcli dev generate-sdks --sdks broker
+```

ellf_broker_sdk-7.0.8/ellf_broker_sdk/__init__.py

@@ -0,0 +1,3 @@
+from .client.full_client import AccessTokenCredential, Client
+
+__all__ = ["Client", "AccessTokenCredential"]

ellf_broker_sdk-7.0.8/ellf_broker_sdk/about.json

@@ -0,0 +1,6 @@
+{
+  "title": "Ellf Broker SDK",
+  "name": "ellf-broker-sdk",
+  "version": "7.0.8",
+  "summary": "Ellf Broker SDK"
+}

ellf_broker_sdk-7.0.8/ellf_broker_sdk/about.py

@@ -0,0 +1,23 @@
+"""
+This file originates in shared/shared_about.py, and
+is copied verbatim into the different projects. It should
+not be edited directly, edit the shared version and replicate
+it everywhere.
+"""
+
+import json
+from pathlib import Path
+
+PWD = Path(__file__).parent
+
+
+with (PWD / "about.json").open("r", encoding="utf8") as file_:
+    _about_data = json.load(file_)
+
+
+__name__ = _about_data["name"]
+__title__ = _about_data.get("title", "")
+__description__ = _about_data.get("description", "")
+__summary__ = _about_data.get("summary", "")
+__version__ = _about_data["version"]
+__prog__ = _about_data.get("prog", "")

ellf_broker_sdk-7.0.8/ellf_broker_sdk/client/base.py

@@ -0,0 +1,256 @@
+import json
+from typing import IO, Any, Iterator, TypeVar
+
+import httpx
+from pydantic import BaseModel, TypeAdapter
+
+from .. import errors
+
+_ReturnModelT = TypeVar("_ReturnModelT")
+
+
+class RawIteratorIO(IO[bytes]):
+    """
+    Creates a IO[bytes] from a bytes iterator.
+
+    Adapted from https://stackoverflow.com/a/12604375/7426717
+    """
+
+    def __init__(self, it: Iterator[bytes]) -> None:
+        self._it = it
+        self._memory = b""
+
+        super().__init__()
+
+    def _read1(self, n: int | None = None) -> bytes:
+        while not self._memory:
+            try:
+                next_memory = next(self._it)
+            except StopIteration:
+                break
+            else:
+                self._memory = next_memory
+
+        chunk = self._memory[:n]
+        self._memory = self._memory[len(chunk) :]
+
+        return chunk
+
+    def read(self, n: int | None = None) -> bytes:
+        chunks: list[bytes] = []
+        if n is None or n < 0:
+            while True:
+                m = self._read1()
+                if not m:
+                    break
+                chunks.append(m)
+        else:
+            while n > 0:
+                m = self._read1(n)
+                if not m:
+                    break
+                n -= len(m)
+                chunks.append(m)
+        return b"".join(chunks)
+
+
+CHUNK_SIZE = 1024
+
+
+class BaseClient:
+    name: str
+    path: str
+    _sync_client: httpx.Client
+    _async_client: httpx.AsyncClient
+
+    def __init__(
+        self, sync_client: httpx.Client, async_client: httpx.AsyncClient, path: str
+    ) -> None:
+        self._sync_client = sync_client
+        self._async_client = async_client
+        if path.startswith("/"):
+            path = path[1:]
+        self.name = path.split("/")[-1]
+        self.path = f"/{path}"
+
+    def request(
+        self,
+        method: str,
+        *,
+        endpoint: str,
+        params: None | BaseModel | dict[str, Any] = None,
+        data: None | BaseModel | dict[str, Any] = None,
+        files: dict[str, Any] | None = None,
+        headers: dict[str, str] | None = None,
+        page: int | None = None,
+        size: int | None = None,
+        params_model: type[BaseModel] | None = None,
+        body_model: type[BaseModel] | None = None,
+        return_model: type[_ReturnModelT] | None = None,
+        stream: bool = False,
+        timeout: (
+            httpx._types.TimeoutTypes | httpx._client.UseClientDefault
+        ) = httpx.USE_CLIENT_DEFAULT,
+    ) -> _ReturnModelT | IO[bytes] | None:
+        req = self._get_validated_request(
+            method=method,
+            endpoint=endpoint,
+            params=params,
+            data=data,
+            files=files,
+            params_model=params_model,
+            body_model=body_model,
+            headers=headers,
+            page=page,
+            size=size,
+            timeout=timeout,
+        )
+        response = self._sync_client.send(req, stream=stream)
+        self._raise_and_handle_errors(response)
+        content = None
+        if (
+            response.headers.get("content-type") == "application/json"
+            and response.status_code != 204
+        ):
+            content = response.json()
+        if return_model is not None:
+            return self._convert_response_to_model(content, return_model)
+        return RawIteratorIO(response.iter_bytes(chunk_size=CHUNK_SIZE))
+
+    async def request_async(
+        self,
+        method: str,
+        *,
+        endpoint: str,
+        params: None | BaseModel | dict[str, Any] = None,
+        data: None | BaseModel | dict[str, Any] = None,
+        files: dict[str, Any] | None = None,
+        headers: dict[str, str] | None = None,
+        page: int | None = None,
+        size: int | None = None,
+        params_model: type[BaseModel] | None = None,
+        body_model: type[BaseModel] | None = None,
+        return_model: type[_ReturnModelT] | None = None,
+        stream: bool = False,
+        timeout: (
+            httpx._types.TimeoutTypes | httpx._client.UseClientDefault
+        ) = httpx.USE_CLIENT_DEFAULT,
+    ) -> _ReturnModelT | IO[bytes] | None:
+        req = self._get_validated_request(
+            method=method,
+            endpoint=endpoint,
+            params=params,
+            data=data,
+            files=files,
+            params_model=params_model,
+            body_model=body_model,
+            headers=headers,
+            page=page,
+            size=size,
+            timeout=timeout,
+        )
+        response = await self._async_client.send(req, stream=stream)
+        self._raise_and_handle_errors(response)
+        content = None
+        if (
+            response.headers.get("content-type") == "application/json"
+            and response.status_code != 204
+        ):
+            content = response.json()
+        if return_model is not None:
+            return self._convert_response_to_model(content, return_model)
+        return RawIteratorIO(response.iter_bytes(CHUNK_SIZE))
+
+    def _get_validated_request(
+        self,
+        *,
+        method: str,
+        endpoint: str,
+        params: None | BaseModel | dict[str, Any] = None,
+        data: None | BaseModel | dict[str, Any] = None,
+        files: dict[str, Any] | None = None,
+        params_model: type[BaseModel] | None = None,
+        body_model: type[BaseModel] | None = None,
+        headers: dict[str, str] | None = None,
+        page: int | None = None,
+        size: int | None = None,
+        timeout: (
+            httpx._types.TimeoutTypes | httpx._client.UseClientDefault
+        ) = httpx.USE_CLIENT_DEFAULT,
+    ) -> httpx.Request:
+        url = self.path
+        if endpoint:
+            if endpoint.startswith("/"):
+                url = endpoint
+            else:
+                url = f"{self.path}/{endpoint}"
+        local_params: dict[str, Any] = {}
+        if page is not None:
+            local_params["page"] = page
+        if size is not None:
+            local_params["size"] = size
+        if params is not None:
+            if isinstance(params, dict):
+                params.update(local_params)
+                if params_model is not None:
+                    # if params_model is passed, use Pydantic to validate params
+                    params = params_model(**params)
+                    params = params.model_dump()
+            elif isinstance(params, BaseModel):
+                params = params.model_dump(exclude_none=True, exclude_defaults=True)
+                params.update(local_params)
+        content = None
+        if data is not None:
+            if isinstance(data, dict):
+                if body_model is not None:
+                    body = body_model(**data)
+                    content = body.model_dump_json()
+                else:
+                    content = json.dumps(data)
+            elif isinstance(data, BaseModel):
+                content = data.model_dump_json()
+        req = self._sync_client.build_request(
+            method,
+            url,
+            headers=headers,
+            params=params,
+            content=content,
+            files=files,
+            timeout=timeout,
+        )
+        return req
+
+    def _raise_and_handle_errors(self, response: httpx.Response) -> None:
+        try:
+            response.raise_for_status()
+        except httpx.HTTPStatusError as e:
+            status = e.response.status_code
+            error_type: str | None = None
+            sdk_error: type[errors.BrokerError] | None = None
+            detail_message: str | None = None
+            try:
+                error_data = e.response.json()
+            except json.JSONDecodeError:
+                error_data = {}
+            detail = error_data.get("detail")
+            if detail and isinstance(detail, dict):
+                error_type = detail.get("type")
+                detail_message = detail.get("message", json.dumps(detail))
+            else:
+                detail_message = json.dumps(error_data)
+            if error_type:
+                sdk_error = getattr(errors, error_type, None)
+            if sdk_error:
+                raise sdk_error(detail=detail_message or "") from e
+            elif status == 401:
+                raise errors.AuthError(detail=detail_message or "") from e
+            else:
+                raise
+
+    def _convert_response_to_model(
+        self, data: Any, return_model: type[_ReturnModelT]
+    ) -> _ReturnModelT | None:
+        """Convert the JSON response to the correct Pydantic model"""
+        if data is None:
+            return data
+        return TypeAdapter(return_model).validate_python(data)

ellf_broker_sdk-7.0.8/ellf_broker_sdk/client/check.py

@@ -0,0 +1,35 @@
+from typing import cast
+
+from ..models import (
+    CheckProgressRequest,
+    CheckProgressResponse,
+    CheckStartRequest,
+    CheckStartResponse,
+)
+from .base import BaseClient
+
+
+class Check(BaseClient):
+    def start(self, body: CheckStartRequest) -> CheckStartResponse:
+        res = self.request(
+            "POST", endpoint="start", data=body, return_model=CheckStartResponse
+        )
+        return cast(CheckStartResponse, res)
+
+    async def start_async(self, body: CheckStartRequest) -> CheckStartResponse:
+        res = await self.request_async(
+            "POST", endpoint="start", data=body, return_model=CheckStartResponse
+        )
+        return cast(CheckStartResponse, res)
+
+    def progress(self, body: CheckProgressRequest) -> CheckProgressResponse:
+        res = self.request(
+            "POST", endpoint="progress", data=body, return_model=CheckProgressResponse
+        )
+        return cast(CheckProgressResponse, res)
+
+    async def progress_async(self, body: CheckProgressRequest) -> CheckProgressResponse:
+        res = await self.request_async(
+            "POST", endpoint="progress", data=body, return_model=CheckProgressResponse
+        )
+        return cast(CheckProgressResponse, res)

ellf_broker_sdk-7.0.8/ellf_broker_sdk/client/data.py

@@ -0,0 +1,105 @@
+from typing import Any, Literal, cast
+
+from ..models import Dataset, DatasetCreating, PageDatasetExample
+from ..ty import Page
+from .base import BaseClient
+
+
+class Data(BaseClient):
+    def all(
+        self,
+        session: bool | None = None,
+        *,
+        page: int | None = None,
+        size: int | None = None,
+    ) -> Page[str]:
+        res = self.request(
+            "GET",
+            endpoint="all",
+            data={"session": session},
+            page=page,
+            size=size,
+            return_model=Page[str],
+        )
+        return cast(Page[str], res)
+
+    async def all_async(
+        self,
+        session: bool | None = None,
+        *,
+        page: int | None = None,
+        size: int | None = None,
+    ) -> Page[str]:
+        res = await self.request_async(
+            "GET",
+            endpoint="all",
+            data={"session": session},
+            page=page,
+            size=size,
+            return_model=Page[str],
+        )
+        return cast(Page[str], res)
+
+    def create(self, body: DatasetCreating) -> Dataset:
+        res = self.request("POST", endpoint="create", data=body, return_model=Dataset)
+        return cast(Dataset, res)
+
+    async def create_async(self, body: DatasetCreating) -> Dataset:
+        res = await self.request_async(
+            "POST", endpoint="create", data=body, return_model=Dataset
+        )
+        return cast(Dataset, res)
+
+    def read_examples(
+        self,
+        datasets: list[str],
+        *,
+        page: int | None = None,
+        size: int | None = None,
+        order: Literal["asc", "desc"] | None = "asc",
+    ) -> PageDatasetExample:
+        params: dict[str, Any] = {}
+        if page is not None:
+            params["page"] = page
+        if size is not None:
+            params["size"] = size
+        if order is not None:
+            params["order"] = order
+
+        data = {"datasets": datasets}
+
+        res = self.request(
+            "POST",
+            endpoint="read-examples",
+            params=params,
+            data=data,
+            return_model=PageDatasetExample,
+        )
+        return cast(PageDatasetExample, res)
+
+    async def read_examples_async(
+        self,
+        datasets: list[str],
+        *,
+        page: int | None = None,
+        size: int | None = None,
+        order: Literal["asc", "desc"] | None = "asc",
+    ) -> PageDatasetExample:
+        params: dict[str, Any] = {}
+        if page is not None:
+            params["page"] = page
+        if size is not None:
+            params["size"] = size
+        if order is not None:
+            params["order"] = order
+
+        data = {"datasets": datasets}
+
+        res = await self.request_async(
+            "POST",
+            endpoint="read-examples",
+            params=params,
+            data=data,
+            return_model=PageDatasetExample,
+        )
+        return cast(PageDatasetExample, res)

ellf_broker_sdk-7.0.8/ellf_broker_sdk/client/envs.py

@@ -0,0 +1,79 @@
+from typing import Any, cast
+
+from pydantic import BaseModel, Field
+
+from ..models import RecipesMeta
+from .base import BaseClient
+
+
+class BuildImageResponse(BaseModel):
+    """Response from the build-image endpoint."""
+
+    image_uri: str = Field(..., title="Image Uri")
+    meta: list[Any] = Field(..., title="Meta")
+
+
+class Envs(BaseClient):
+    def get_recipes_meta(self, package_name: str, package_version: str) -> RecipesMeta:
+        res = self.request(
+            "GET",
+            endpoint=f"{package_name}/{package_version}/meta",
+            return_model=RecipesMeta,
+        )
+        return cast(RecipesMeta, res)
+
+    async def get_recipes_meta_async(
+        self, package_name: str, package_version: str
+    ) -> RecipesMeta:
+        res = await self.request_async(
+            "GET",
+            endpoint=f"{package_name}/{package_version}/meta",
+            return_model=RecipesMeta,
+        )
+        return cast(RecipesMeta, res)
+
+    def build_image(
+        self,
+        package_name: str,
+        package_version: str,
+        base_image: str,
+        tar_data_b64: str,
+        meta_json: str,
+    ) -> BuildImageResponse:
+        """Build an ad-hoc container image for a recipe package."""
+        res = self.request(
+            "POST",
+            endpoint="build-image",
+            data={
+                "package_name": package_name,
+                "package_version": package_version,
+                "base_image": base_image,
+                "tar_data": tar_data_b64,
+                "meta_json": meta_json,
+            },
+            return_model=BuildImageResponse,
+        )
+        return cast(BuildImageResponse, res)
+
+    async def build_image_async(
+        self,
+        package_name: str,
+        package_version: str,
+        base_image: str,
+        tar_data_b64: str,
+        meta_json: str,
+    ) -> BuildImageResponse:
+        """Build an ad-hoc container image for a recipe package (async)."""
+        res = await self.request_async(
+            "POST",
+            endpoint="build-image",
+            data={
+                "package_name": package_name,
+                "package_version": package_version,
+                "base_image": base_image,
+                "tar_data": tar_data_b64,
+                "meta_json": meta_json,
+            },
+            return_model=BuildImageResponse,
+        )
+        return cast(BuildImageResponse, res)

ellf_broker_sdk-7.0.8/ellf_broker_sdk/client/files.py

@@ -0,0 +1,139 @@
+from typing import IO, cast
+
+import httpx
+
+from ..models import (
+    Copying,
+    CopyPlan,
+    Deleting,
+    Downloading,
+    FileStats,
+    Listing,
+    PathList,
+    RsyncPlan,
+    SignedUrl,
+    Statting,
+)
+from .base import BaseClient
+
+
+class Files(BaseClient):
+    def get_signed_url(self, *, path: str, method: str) -> SignedUrl:
+        res = self.request(
+            "POST",
+            endpoint="signed-url",
+            return_model=SignedUrl,
+            data={"path": path, "method": method},
+        )
+        return cast(SignedUrl, res)
+
+    async def get_signed_url_async(self, *, path: str, method: str) -> SignedUrl:
+        res = await self.request_async(
+            "POST",
+            endpoint="signed-url",
+            return_model=SignedUrl,
+            data={"path": path, "method": method},
+        )
+        return cast(SignedUrl, res)
+
+    def upload(
+        self, file: IO[bytes], *, dest: str, overwrite: bool, make_dirs: bool
+    ) -> None:
+        self.request(
+            "POST",
+            endpoint="upload",
+            files={"file": file},
+            params={"dest": dest, "overwrite": overwrite, "make_dirs": make_dirs},
+            return_model=SignedUrl,
+            timeout=httpx.Timeout(5.0, read=300.0, write=300.0),
+        )
+
+    async def upload_async(
+        self, file: IO[bytes], *, dest: str, overwrite: bool, make_dirs: bool
+    ) -> None:
+        await self.request_async(
+            "POST",
+            endpoint="upload",
+            files={"file": file},
+            params={"dest": dest, "overwrite": overwrite, "make_dirs": make_dirs},
+            return_model=None,
+            timeout=httpx.Timeout(5.0, read=300.0, write=300.0),
+        )
+
+    def copy(self, body: Copying) -> None:
+        self.request("POST", endpoint="copy", data=body, return_model=None)
+
+    async def copy_async(self, body: Copying) -> None:
+        await self.request_async("POST", endpoint="copy", data=body, return_model=None)
+
+    def delete(self, body: Deleting) -> None:
+        self.request("POST", endpoint="delete", data=body, return_model=None)
+
+    async def delete_async(self, body: Deleting) -> None:
+        await self.request_async(
+            "POST", endpoint="delete", data=body, return_model=None
+        )
+
+    def stat(self, body: Statting) -> FileStats:
+        res = self.request("POST", endpoint="stat", data=body, return_model=FileStats)
+        return cast(FileStats, res)
+
+    async def stat_async(self, body: Statting) -> FileStats:
+        res = await self.request_async(
+            "POST", endpoint="stat", data=body, return_model=FileStats
+        )
+        return cast(FileStats, res)
+
+    def list_dir(self, body: Listing) -> PathList:
+        res = self.request(
+            "POST",
+            endpoint="list-dir",
+            data=body,
+            return_model=PathList,
+            timeout=httpx.Timeout(5.0, read=20.0),
+        )
+        return cast(PathList, res)
+
+    async def list_dir_async(self, body: Listing) -> PathList:
+        res = await self.request_async(
+            "POST",
+            endpoint="list-dir",
+            data=body,
+            return_model=PathList,
+            timeout=httpx.Timeout(5.0, read=20.0),
+        )
+        return cast(PathList, res)
+
+    def plan_directory_copy(self, body: Copying) -> CopyPlan:
+        res = self.request(
+            "POST", endpoint="plan-directory-copy", data=body, return_model=CopyPlan
+        )
+        return cast(CopyPlan, res)
+
+    async def plan_directory_copy_async(self, body: Copying) -> CopyPlan:
+        res = await self.request_async(
+            "POST", endpoint="plan-directory-copy", data=body, return_model=CopyPlan
+        )
+        return cast(CopyPlan, res)
+
+    def plan_directory_rsync(self, body: Copying) -> RsyncPlan:
+        res = self.request(
+            "POST", endpoint="plan-directory-rsync", data=body, return_model=RsyncPlan
+        )
+        return cast(RsyncPlan, res)
+
+    async def plan_directory_rsync_async(self, body: Copying) -> RsyncPlan:
+        res = await self.request_async(
+            "POST", endpoint="plan-directory-rsync", data=body, return_model=RsyncPlan
+        )
+        return cast(RsyncPlan, res)
+
+    def download(self, body: Downloading) -> IO[bytes]:
+        res = self.request("POST", endpoint="download", data=body, stream=True)
+        return cast(IO[bytes], res)
+
+    async def download_async(self, body: Downloading) -> IO[bytes]:
+        res = await self.request_async(
+            "POST", endpoint="download", data=body, stream=True
+        )
+        return cast(IO[bytes], res)