flaxcloud 0.3.0__py3-none-any.whl

flaxcloud/__init__.py

@@ -0,0 +1,59 @@
+"""Flax Cloud — Python SDK and CLI for isolated cloud sandboxes.
+
+```python
+from flaxcloud import FlaxClient
+
+flax = FlaxClient()                       # reads FLAX_API_KEY
+with flax.create_sandbox(template="python") as sb:
+    print(sb.run("python3 -c 'print(6*7)'").stdout)
+```
+"""
+
+from __future__ import annotations
+
+from .client import FlaxClient
+from .errors import (
+    FlaxAuthError,
+    FlaxBadRequestError,
+    FlaxConflictError,
+    FlaxConnectionError,
+    FlaxError,
+    FlaxNotFoundError,
+    FlaxQuotaError,
+    FlaxServerError,
+)
+from .models import Command, FileEntry, PreviewLink, SessionResult, StartupStatus
+from .sandbox import Sandbox, Session
+
+__version__ = "0.3.0"
+
+
+def __getattr__(name: str):
+    # Lazily expose the async client so importing the package doesn't require asyncio setup.
+    if name in ("AsyncFlaxClient", "AsyncSandbox"):
+        from . import aio
+
+        return getattr(aio, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "FlaxClient",
+    "Sandbox",
+    "Session",
+    "SessionResult",
+    "AsyncFlaxClient",
+    "AsyncSandbox",
+    "Command",
+    "FileEntry",
+    "PreviewLink",
+    "StartupStatus",
+    "FlaxError",
+    "FlaxAuthError",
+    "FlaxNotFoundError",
+    "FlaxQuotaError",
+    "FlaxConflictError",
+    "FlaxBadRequestError",
+    "FlaxServerError",
+    "FlaxConnectionError",
+]

flaxcloud/aio.py

@@ -0,0 +1,440 @@
+"""Async client for Flax Cloud — `AsyncFlaxClient` + `AsyncSandbox`.
+
+Mirrors the sync `FlaxClient`/`Sandbox` API using `httpx.AsyncClient` and `await`.
+
+```python
+import asyncio
+from flaxcloud.aio import AsyncFlaxClient
+
+async def main():
+    async with AsyncFlaxClient() as flax:
+        sb = await flax.create_sandbox(template="python")
+        out = await sb.run("python3 -c 'print(6*7)'")
+        print(out.stdout)
+        await sb.destroy()
+
+asyncio.run(main())
+```
+"""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import os
+import random
+import shlex
+from typing import Any, Optional, Union
+
+import httpx
+
+from .client import (
+    _IDEMPOTENT,
+    _RETRY_STATUSES,
+    DEFAULT_BASE_URL,
+    _user_agent,
+)
+from .errors import FlaxAuthError, FlaxConnectionError, error_from_response
+from .models import Command, FileEntry, PreviewLink, SessionResult, StartupStatus
+from .sandbox import WORKSPACE
+
+
+class AsyncFlaxClient:
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        *,
+        base_url: Optional[str] = None,
+        timeout: float = 60.0,
+        max_retries: int = 2,
+        http_client: Optional[httpx.AsyncClient] = None,
+    ) -> None:
+        self._api_key = api_key or os.environ.get("FLAX_API_KEY")
+        if not self._api_key:
+            raise FlaxAuthError(
+                "no API key: pass api_key=... or set the FLAX_API_KEY environment variable"
+            )
+        base_url = base_url or os.environ.get("FLAX_BASE_URL") or DEFAULT_BASE_URL
+        self._max_retries = max(0, int(max_retries))
+        self._http = http_client or httpx.AsyncClient(base_url=base_url.rstrip("/"), timeout=timeout)
+        self._owns_client = http_client is None
+        self.base_url = base_url.rstrip("/")
+
+    async def request(self, method: str, path: str, **kwargs: Any) -> Any:
+        headers = dict(kwargs.pop("headers", {}))
+        headers["Authorization"] = f"Bearer {self._api_key}"
+        headers.setdefault("User-Agent", _user_agent())
+        method = method.upper()
+
+        attempt = 0
+        while True:
+            try:
+                resp = await self._http.request(method, path, headers=headers, **kwargs)
+            except httpx.ConnectError as exc:
+                if attempt < self._max_retries:
+                    await self._sleep_backoff(attempt)
+                    attempt += 1
+                    continue
+                raise FlaxConnectionError(f"could not reach Flax Cloud: {exc}") from exc
+            except httpx.HTTPError as exc:
+                if method in _IDEMPOTENT and attempt < self._max_retries:
+                    await self._sleep_backoff(attempt)
+                    attempt += 1
+                    continue
+                raise FlaxConnectionError(f"could not reach Flax Cloud: {exc}") from exc
+
+            if (
+                resp.status_code in _RETRY_STATUSES
+                and method in _IDEMPOTENT
+                and attempt < self._max_retries
+            ):
+                await self._sleep_backoff(attempt, resp.headers.get("Retry-After"))
+                attempt += 1
+                continue
+            return self._parse(resp)
+
+    @staticmethod
+    async def _sleep_backoff(attempt: int, retry_after: Optional[str] = None) -> None:
+        if retry_after:
+            try:
+                await asyncio.sleep(min(float(retry_after), 30.0))
+                return
+            except (TypeError, ValueError):
+                pass
+        await asyncio.sleep(min(0.25 * (2 ** attempt), 5.0) * (0.5 + random.random()))
+
+    @staticmethod
+    def _parse(resp: httpx.Response) -> Any:
+        if resp.status_code >= 400:
+            code = None
+            message = resp.text or resp.reason_phrase
+            try:
+                payload = resp.json()
+                if isinstance(payload, dict) and isinstance(payload.get("error"), dict):
+                    code = payload["error"].get("code")
+                    message = payload["error"].get("message", message)
+            except ValueError:
+                pass
+            raise error_from_response(resp.status_code, code, message)
+        if resp.status_code == 204 or not resp.content:
+            return None
+        return resp.json()
+
+    # ------------------------------- sandboxes -------------------------------
+    async def create_sandbox(
+        self,
+        template: str = "blank",
+        *,
+        memory_mb: Optional[int] = None,
+        timeout_seconds: Optional[int] = None,
+        network_mode: Optional[str] = None,
+        env: Optional[dict] = None,
+        startup_command: Optional[str] = None,
+        restore_from: Optional[str] = None,
+        metadata: Optional[dict] = None,
+    ) -> "AsyncSandbox":
+        body: dict[str, Any] = {"template": template}
+        if memory_mb is not None:
+            body["memory_mb"] = memory_mb
+        if timeout_seconds is not None:
+            body["timeout_seconds"] = timeout_seconds
+        if network_mode is not None:
+            body["network_mode"] = network_mode
+        if env is not None:
+            body["env"] = env
+        if startup_command is not None:
+            body["startup_command"] = startup_command
+        if restore_from is not None:
+            body["restore_from"] = restore_from
+        if metadata is not None:
+            body["metadata"] = metadata
+        return AsyncSandbox(self, await self.request("POST", "/v1/sandboxes", json=body))
+
+    async def get_sandbox(self, sandbox_id: str) -> "AsyncSandbox":
+        return AsyncSandbox(self, await self.request("GET", f"/v1/sandboxes/{sandbox_id}"))
+
+    async def list_sandboxes(self) -> list["AsyncSandbox"]:
+        data = await self.request("GET", "/v1/sandboxes")
+        return [AsyncSandbox(self, d) for d in data.get("data", [])]
+
+    async def me(self) -> dict:
+        return await self.request("GET", "/v1/me")
+
+    async def usage(self) -> dict:
+        return await self.request("GET", "/v1/usage")
+
+    # ------------------------------- lifecycle -------------------------------
+    async def aclose(self) -> None:
+        if self._owns_client:
+            await self._http.aclose()
+
+    async def __aenter__(self) -> "AsyncFlaxClient":
+        return self
+
+    async def __aexit__(self, *exc: object) -> None:
+        await self.aclose()
+
+
+class AsyncSandbox:
+    def __init__(self, client: AsyncFlaxClient, data: dict) -> None:
+        self._client = client
+        self._data = data
+
+    @property
+    def id(self) -> str:
+        return self._data["id"]
+
+    @property
+    def status(self) -> str:
+        return self._data.get("status", "")
+
+    @property
+    def env(self) -> dict:
+        return dict(self._data.get("env") or {})
+
+    def __repr__(self) -> str:  # pragma: no cover - cosmetic
+        return f"<AsyncSandbox {self.id} status={self.status!r}>"
+
+    async def refresh(self) -> "AsyncSandbox":
+        self._data = await self._client.request("GET", f"/v1/sandboxes/{self.id}")
+        return self
+
+    async def stop(self) -> "AsyncSandbox":
+        self._data = await self._client.request("POST", f"/v1/sandboxes/{self.id}/stop")
+        return self
+
+    async def resume(self) -> "AsyncSandbox":
+        self._data = await self._client.request("POST", f"/v1/sandboxes/{self.id}/resume")
+        return self
+
+    async def destroy(self) -> None:
+        await self._client.request("DELETE", f"/v1/sandboxes/{self.id}")
+
+    async def __aenter__(self) -> "AsyncSandbox":
+        return self
+
+    async def __aexit__(self, *exc: object) -> None:
+        try:
+            await self.destroy()
+        except Exception:  # noqa: BLE001
+            pass
+
+    async def run(
+        self,
+        command: str,
+        *,
+        timeout: Optional[int] = None,
+        working_directory: Optional[str] = None,
+        env: Optional[dict] = None,
+        background: bool = False,
+    ) -> Command:
+        body: dict = {"command": command}
+        if timeout is not None:
+            body["timeout_seconds"] = timeout
+        if working_directory is not None:
+            body["working_directory"] = working_directory
+        if env is not None:
+            body["env"] = env
+        if background:
+            body["async"] = True
+        return Command.from_dict(
+            await self._client.request("POST", f"/v1/sandboxes/{self.id}/commands", json=body)
+        )
+
+    async def code_run(
+        self, code: str, *, language: str = "python", timeout: Optional[int] = None,
+        env: Optional[dict] = None, background: bool = False,
+    ) -> Command:
+        from .sandbox import Sandbox
+
+        interp = Sandbox._CODE_INTERPRETERS.get(language.lower())
+        if interp is None:
+            raise ValueError(f"unsupported language {language!r}")
+        encoded = base64.b64encode(code.encode("utf-8")).decode("ascii")
+        return await self.run(
+            f"echo {encoded} | base64 -d | {interp}",
+            timeout=timeout, env=env, background=background,
+        )
+
+    def run_stream(
+        self, command: str, *, timeout: Optional[int] = None, env: Optional[dict] = None
+    ) -> "AsyncCommandStream":
+        body: dict = {"command": command}
+        if timeout is not None:
+            body["timeout_seconds"] = timeout
+        if env is not None:
+            body["env"] = env
+        return AsyncCommandStream(self._client, f"/v1/sandboxes/{self.id}/commands/stream", body)
+
+    async def get_command(self, command_id: str) -> Command:
+        return Command.from_dict(await self._client.request("GET", f"/v1/commands/{command_id}"))
+
+    async def wait(
+        self, command: Union[Command, str], *, interval: float = 1.0, timeout: float = 300.0
+    ) -> Command:
+        command_id = command.id if isinstance(command, Command) else command
+        loop = asyncio.get_event_loop()
+        deadline = loop.time() + timeout
+        current = await self.get_command(command_id)
+        while not current.done:
+            if loop.time() > deadline:
+                raise TimeoutError(f"command {command_id} did not finish within {timeout}s")
+            await asyncio.sleep(interval)
+            current = await self.get_command(command_id)
+        return current
+
+    # files
+    async def upload(self, path: str, content: Union[bytes, str]) -> None:
+        raw = content.encode("utf-8") if isinstance(content, str) else content
+        await self._client.request(
+            "POST", f"/v1/sandboxes/{self.id}/files",
+            json={"path": path, "content_base64": base64.b64encode(raw).decode()},
+        )
+
+    async def download(self, path: str) -> bytes:
+        data = await self._client.request(
+            "GET", f"/v1/sandboxes/{self.id}/files", params={"path": path}
+        )
+        return base64.b64decode(data["content_base64"])
+
+    async def list_files(self, path: str = WORKSPACE) -> list[FileEntry]:
+        data = await self._client.request(
+            "GET", f"/v1/sandboxes/{self.id}/files/list", params={"path": path}
+        )
+        return [FileEntry.from_dict(d) for d in data.get("data", [])]
+
+    async def delete_file(self, path: str) -> None:
+        await self._client.request(
+            "DELETE", f"/v1/sandboxes/{self.id}/files", params={"path": path}
+        )
+
+    async def read_text(self, path: str, encoding: str = "utf-8") -> str:
+        return (await self.download(path)).decode(encoding)
+
+    async def mkdir(self, path: str) -> None:
+        await self.run(f"mkdir -p {shlex.quote(path)}")
+
+    async def move(self, src: str, dst: str) -> None:
+        await self.run(f"mv {shlex.quote(src)} {shlex.quote(dst)}")
+
+    async def exists(self, path: str) -> bool:
+        return (await self.run(f"test -e {shlex.quote(path)}")).exit_code == 0
+
+    # startup
+    async def startup(self) -> StartupStatus:
+        return StartupStatus.from_dict(
+            await self._client.request("GET", f"/v1/sandboxes/{self.id}/startup")
+        )
+
+    async def set_startup(self, command: Optional[str]) -> StartupStatus:
+        return StartupStatus.from_dict(
+            await self._client.request(
+                "PUT", f"/v1/sandboxes/{self.id}/startup", json={"command": command}
+            )
+        )
+
+    async def run_startup(self) -> StartupStatus:
+        return StartupStatus.from_dict(
+            await self._client.request("POST", f"/v1/sandboxes/{self.id}/startup/run")
+        )
+
+    # previews
+    async def create_preview_link(self, port: int) -> PreviewLink:
+        return PreviewLink.from_dict(
+            await self._client.request(
+                "POST", f"/v1/sandboxes/{self.id}/preview-link", json={"port": port}
+            )
+        )
+
+    # sessions
+    async def create_session(self) -> "AsyncSession":
+        data = await self._client.request("POST", f"/v1/sandboxes/{self.id}/sessions")
+        return AsyncSession(self._client, data)
+
+    async def sessions(self) -> list["AsyncSession"]:
+        data = await self._client.request("GET", f"/v1/sandboxes/{self.id}/sessions")
+        return [AsyncSession(self._client, d) for d in data.get("data", [])]
+
+
+class AsyncCommandStream:
+    """Async iterable of live output chunks; `exit_code` is set once the stream ends."""
+
+    def __init__(self, client: AsyncFlaxClient, path: str, body: dict) -> None:
+        self._client = client
+        self._path = path
+        self._body = body
+        self.exit_code: Optional[int] = None
+
+    async def __aiter__(self):
+        from .errors import FlaxError
+
+        headers = {
+            "Authorization": f"Bearer {self._client._api_key}",
+            "User-Agent": _user_agent(),
+            "Accept": "text/event-stream",
+        }
+        async with self._client._http.stream(
+            "POST", self._path, json=self._body, headers=headers
+        ) as resp:
+            if resp.status_code >= 400:
+                await resp.aread()
+                AsyncFlaxClient._parse(resp)
+            error: Optional[str] = None
+            event: Optional[str] = None
+            async for line in resp.aiter_lines():
+                if not line:
+                    event = None
+                    continue
+                if line.startswith("event:"):
+                    event = line[len("event:"):].strip()
+                elif line.startswith("data:"):
+                    data = line[len("data:"):].strip()
+                    if event == "chunk":
+                        yield base64.b64decode(data).decode("utf-8", "replace")
+                    elif event == "error":
+                        error = base64.b64decode(data).decode("utf-8", "replace")
+                    elif event == "end":
+                        try:
+                            self.exit_code = int(data)
+                        except ValueError:
+                            self.exit_code = None
+        if error is not None:
+            raise FlaxError(f"streamed command failed: {error}")
+
+
+class AsyncSession:
+    """Async stateful session; cwd/exported env persist across `run()` calls."""
+
+    def __init__(self, client: AsyncFlaxClient, data: dict) -> None:
+        self._client = client
+        self._data = data
+
+    @property
+    def id(self) -> str:
+        return self._data["id"]
+
+    @property
+    def cwd(self) -> str:
+        return self._data.get("cwd", WORKSPACE)
+
+    async def run(self, command: str, *, timeout: Optional[int] = None) -> SessionResult:
+        body: dict = {"command": command}
+        if timeout is not None:
+            body["timeout_seconds"] = timeout
+        result = SessionResult.from_dict(
+            await self._client.request("POST", f"/v1/sessions/{self.id}/exec", json=body)
+        )
+        self._data["cwd"] = result.cwd
+        return result
+
+    async def delete(self) -> None:
+        await self._client.request("DELETE", f"/v1/sessions/{self.id}")
+
+    async def __aenter__(self) -> "AsyncSession":
+        return self
+
+    async def __aexit__(self, *exc: object) -> None:
+        try:
+            await self.delete()
+        except Exception:  # noqa: BLE001
+            pass