akernel-sdk 0.1.0__py3-none-any.whl

akernel_sdk/_openyuanrong.py

@@ -0,0 +1,320 @@
+# Copyright (c) 2026 Ant Group Corporation.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Internal openYuanrong adapter used by the public AKernel SDK."""
+
+from __future__ import annotations
+
+import atexit
+import json
+import logging
+import os
+import threading
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+import yr
+from yr.runtime_holder import global_runtime
+from yr.sandbox.sandbox import _sanitize_instance_id
+
+from ._addresses import api_endpoint_from_env, gateway_endpoint_from_env
+from ._instance import _SandboxInstance
+from .types import HttpReverseTunnel, Mount, NodeInfo, S3Config
+
+logger = logging.getLogger(__name__)
+
+_NAMESPACE = "akernel"
+_initialized = False
+_init_lock = threading.Lock()
+
+
+def ensure_initialized() -> None:
+    """Initialize the openYuanrong client once for the current process."""
+
+    global _initialized
+    if _initialized:
+        return
+
+    with _init_lock:
+        if _initialized:
+            return
+
+        token = os.environ.get("AKERNEL_TOKEN", "").strip()
+        if not token:
+            raise RuntimeError("AKERNEL_TOKEN is not set")
+
+        endpoint = api_endpoint_from_env()
+        # Keep openYuanrong tuning inside its backend initialization boundary.
+        os.environ.setdefault("YR_HTTP_CONNECTION_NUM", "2")
+        os.environ.setdefault("YR_SERVER_ADDRESS", endpoint.authority())
+        gateway_override = os.environ.get("AKERNEL_GATEWAY_ADDRESS", "").strip()
+        if gateway_override and not os.environ.get("YR_GATEWAY_ADDRESS", "").strip():
+            os.environ["YR_GATEWAY_ADDRESS"] = gateway_override
+
+        yr.init(
+            yr.Config(
+                server_address=endpoint.authority(),
+                auth_token=token,
+                enable_tls=endpoint.use_tls,
+                server_name="akernel-sdk",
+                in_cluster=False,
+                http_ioc_threads_num=8,
+                bypass_datasystem=True,
+            )
+        )
+        atexit.register(yr.finalize)
+        _initialized = True
+
+
+def _validate_positive_int(name: str, value: int, *, allow_zero: bool = False) -> None:
+    if isinstance(value, bool) or not isinstance(value, int):
+        raise TypeError(f"{name} must be an integer")
+    minimum = 0 if allow_zero else 1
+    if value < minimum:
+        comparator = "non-negative" if allow_zero else "greater than 0"
+        raise ValueError(f"{name} must be {comparator}")
+
+
+def _rootfs_json(
+    *, image: str | None, rootfs: S3Config | None, runtime: str
+) -> str | None:
+    if image is not None:
+        return json.dumps(
+            {
+                "runtime": runtime,
+                "type": "image",
+                "readonly": False,
+                "imageurl": image,
+            }
+        )
+    if rootfs is not None:
+        return json.dumps(
+            {
+                "runtime": runtime,
+                "type": "s3",
+                "readonly": False,
+                "s3_config": rootfs.to_dict(),
+            }
+        )
+    return None
+
+
+def build_options(
+    *,
+    image: str | None,
+    rootfs: S3Config | None,
+    runtime: str,
+    cpu: int,
+    memory: int,
+    cpu_limit: int,
+    mem_limit: int,
+    idle_timeout: int,
+    schedule_timeout: int,
+    env: Mapping[str, str] | None,
+    name: str | None,
+    port_forwardings: Sequence[int],
+    mounts: Sequence[Mount],
+    reverse_tunnel: HttpReverseTunnel | None,
+    detached: bool,
+    node_id: str | None,
+) -> Any:
+    """Translate the stable SDK configuration to openYuanrong options."""
+
+    _validate_positive_int("cpu", cpu)
+    _validate_positive_int("memory", memory)
+    _validate_positive_int("cpu_limit", cpu_limit, allow_zero=True)
+    _validate_positive_int("mem_limit", mem_limit, allow_zero=True)
+    _validate_positive_int("idle_timeout", idle_timeout, allow_zero=True)
+    if schedule_timeout != -1:
+        _validate_positive_int("schedule_timeout", schedule_timeout, allow_zero=True)
+    if cpu_limit and cpu_limit < cpu:
+        raise ValueError("cpu_limit must be 0 or greater than or equal to cpu")
+    if mem_limit and mem_limit < memory:
+        raise ValueError("mem_limit must be 0 or greater than or equal to memory")
+
+    options = yr.InvokeOptions()
+    # A Sandbox is driven by one sequential SDK client. Disabling ordered RPC
+    # execution prevents a missing sequence number from stalling later calls.
+    options.need_order = False
+    options.idle_timeout = idle_timeout
+    options.schedule_timeout_ms = (
+        -1 if schedule_timeout == -1 else schedule_timeout * 1000
+    )
+    options.cpu = cpu
+    options.memory = memory
+    options.cpu_limit = cpu_limit
+    options.mem_limit = mem_limit
+    options.namespace = _NAMESPACE
+
+    rootfs_json = _rootfs_json(image=image, rootfs=rootfs, runtime=runtime)
+    if rootfs_json is not None:
+        options.custom_extensions["rootfs"] = rootfs_json
+    if detached:
+        options.custom_extensions["lifecycle"] = "detached"
+    if env:
+        options.env_vars = dict(env)
+    if name:
+        options.name = name
+    if mounts:
+        options.custom_extensions["mounts"] = json.dumps(
+            [mount.to_dict() for mount in mounts]
+        )
+
+    forwarded = list(port_forwardings)
+    if reverse_tunnel is not None:
+        forwarded.append(reverse_tunnel.reverse_port)
+    if forwarded:
+        from yr.config import PortForwarding
+
+        options.port_forwardings = [
+            PortForwarding(port=port, protocol="TCP") for port in forwarded
+        ]
+
+    if node_id is not None:
+        if not isinstance(node_id, str) or not node_id.strip():
+            raise ValueError("node_id must be a non-empty string")
+        operator = yr.LabelOperator(
+            yr.OperatorType.LABEL_IN,
+            "NODE_ID",
+            [node_id.strip()],
+        )
+        options.schedule_affinities = [
+            yr.Affinity(
+                yr.AffinityKind.RESOURCE,
+                yr.AffinityType.REQUIRED,
+                [operator],
+            )
+        ]
+
+    return options
+
+
+def create_instance(options: Any, *, cwd: str | None) -> Any:
+    instance_class: Any = _SandboxInstance
+    handle = instance_class.options(options).invoke(cwd=cwd)
+    yr.get(handle.ping.invoke())
+    return handle
+
+
+def real_instance_id(handle: Any) -> str:
+    return global_runtime.get_runtime().get_real_instance_id(handle.instance_id)
+
+
+def terminate_instance(handle: Any) -> None:
+    handle.terminate()
+
+
+def delete_named_instance(name: str) -> None:
+    ensure_initialized()
+    handle = yr.get_instance(name, namespace=_NAMESPACE)
+    handle.terminate(is_sync=True)
+
+
+def ping_instance(handle: Any) -> bool:
+    try:
+        yr.get(handle.ping.invoke())
+        return True
+    except Exception:
+        return False
+
+
+def instance_state(handle: Any) -> str:
+    result = yr.get(handle.get_info.invoke())
+    return result.get("state", "unknown")
+
+
+def start_reverse_tunnel(
+    handle: Any,
+    tunnel: HttpReverseTunnel,
+    *,
+    name: str | None,
+) -> Any:
+    """Start both ends of an HTTP reverse tunnel and return its client."""
+
+    yr.get(
+        handle.start_tunnel_server.invoke(
+            tunnel.reverse_port,
+            tunnel.listen_port,
+        )
+    )
+    gateway = gateway_endpoint_from_env()
+    instance_id = real_instance_id(handle)
+    safe_id = _sanitize_instance_id(instance_id)
+    websocket_url = (
+        f"{gateway.websocket_scheme}://"
+        f"{gateway.authority(omit_default_port=True)}/"
+        f"{safe_id}/{tunnel.reverse_port}"
+    )
+
+    from yr.sandbox.tunnel_client import TunnelClient
+
+    client = TunnelClient(tunnel.target)
+    logger.info(
+        "Starting reverse tunnel: sandbox_id=%s name=%s url=%s target=%s",
+        safe_id,
+        name or "",
+        websocket_url,
+        tunnel.target,
+    )
+    if client.start(websocket_url, timeout=float(tunnel.connect_timeout)):
+        return client
+
+    client.stop()
+    raise RuntimeError(
+        "reverse tunnel connection timed out after "
+        f"{float(tunnel.connect_timeout):.1f}s: sandbox_id={safe_id} "
+        f"name={name or ''} tunnel_url={websocket_url}"
+    )
+
+
+def _to_float_mapping(value: Any) -> dict[str, float]:
+    if not isinstance(value, Mapping):
+        return {}
+    result: dict[str, float] = {}
+    for key, raw in value.items():
+        scalar = raw
+        if isinstance(raw, Mapping):
+            scalar = raw.get("value", raw.get("scalar", raw))
+            if isinstance(scalar, Mapping):
+                scalar = scalar.get("value", 0)
+        try:
+            result[str(key)] = float(scalar)
+        except (TypeError, ValueError):
+            continue
+    return result
+
+
+def _to_node_info(value: Any) -> NodeInfo:
+    if not isinstance(value, Mapping):
+        raise TypeError(f"unexpected openYuanrong resource entry: {value!r}")
+    labels = value.get("labels", value.get("nodeLabels", {}))
+    return NodeInfo(
+        id=str(value.get("id", "")),
+        status=int(value.get("status", 0)),
+        capacity=_to_float_mapping(value.get("capacity", {})),
+        allocatable=_to_float_mapping(value.get("allocatable", {})),
+        labels=dict(labels) if isinstance(labels, Mapping) else {},
+    )
+
+
+def resources() -> list[NodeInfo]:
+    """Return cluster resources through stable AKernel value types."""
+
+    ensure_initialized()
+    raw = yr.resources()
+    if raw is None:
+        return []
+    if not isinstance(raw, Sequence) or isinstance(raw, (str, bytes)):
+        raise TypeError("openYuanrong resources() returned a non-list value")
+    return [_to_node_info(item) for item in raw]

akernel_sdk/_pty_transport.py

@@ -0,0 +1,304 @@
+# Copyright (c) 2026 Ant Group Corporation.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Thread-backed WebSocket transport for the synchronous PTY API."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import threading
+from collections.abc import Callable, Sequence
+from concurrent.futures import Future
+from typing import Any
+from urllib.parse import urlencode
+
+from ._addresses import Endpoint
+
+_PROTOCOL = "sandbox.pty.v1"
+_PROTOCOL_VERSION = 1
+_WRITE_TIMEOUT = 10.0
+
+
+class _PtyTransportError(RuntimeError):
+    """Internal transport failure converted to PtyError by the public API."""
+
+
+def _build_pty_uri(
+    endpoint: Endpoint,
+    *,
+    instance_id: str,
+    token: str,
+    command: Sequence[str],
+    rows: int,
+    cols: int,
+) -> str:
+    parameters: list[tuple[str, str]] = [
+        ("instance", instance_id),
+        ("tenant_id", "default"),
+        ("token", token),
+        ("tty", "true"),
+        ("rows", str(rows)),
+        ("cols", str(cols)),
+        ("protocol", _PROTOCOL),
+    ]
+    parameters.extend(("command", argument) for argument in command)
+    return (
+        f"{endpoint.websocket_scheme}://"
+        f"{endpoint.authority(omit_default_port=True)}/terminal/ws?"
+        f"{urlencode(parameters)}"
+    )
+
+
+class _PtyConnection:
+    def __init__(
+        self,
+        uri: str,
+        *,
+        ssl_context: Any,
+        rows: int,
+        cols: int,
+        on_data: Callable[[bytes], None] | None,
+        on_done: Callable[[], None],
+    ) -> None:
+        self._uri = uri
+        self._ssl_context = ssl_context
+        self._rows = rows
+        self._cols = cols
+        self._on_data = on_data
+        self._on_done = on_done
+
+        self._ready = threading.Event()
+        self._done = threading.Event()
+        self._state_lock = threading.Lock()
+        self._thread = threading.Thread(
+            target=self._thread_main,
+            name="akernel-pty",
+            daemon=True,
+        )
+        self._loop: asyncio.AbstractEventLoop | None = None
+        self._run_task: asyncio.Task[None] | None = None
+        self._websocket: Any = None
+        self._session_id: str | None = None
+        self._exit_code: int | None = None
+        self._error: BaseException | None = None
+        self._closed_by_client = False
+
+    @property
+    def session_id(self) -> str:
+        if self._session_id is None:
+            raise _PtyTransportError("PTY session has not started")
+        return self._session_id
+
+    @property
+    def exit_code(self) -> int | None:
+        return self._exit_code
+
+    @property
+    def done(self) -> bool:
+        return self._done.is_set()
+
+    def start(self, timeout: float) -> None:
+        self._thread.start()
+        if not self._ready.wait(timeout):
+            self.close()
+            raise TimeoutError(f"PTY did not start within {timeout:g} seconds")
+        self._raise_error()
+        if self._session_id is None:
+            raise _PtyTransportError("PTY connection closed before it started")
+
+    def send_stdin(self, data: bytes) -> None:
+        self._send(data)
+
+    def close_stdin(self) -> None:
+        self._send("STDIN_EOF")
+
+    def resize(self, *, rows: int, cols: int) -> None:
+        self._send(f"RESIZE:{cols}:{rows}")
+
+    def wait(self, timeout: float | None) -> int:
+        if not self._done.wait(timeout):
+            rendered = (
+                "the requested timeout"
+                if timeout is None
+                else f"{timeout:g} seconds"
+            )
+            raise TimeoutError(f"PTY did not exit within {rendered}")
+        self._raise_error()
+        if self._exit_code is None:
+            if self._closed_by_client:
+                raise _PtyTransportError("PTY session was closed by the client")
+            raise _PtyTransportError("PTY connection closed without an exit status")
+        return self._exit_code
+
+    def close(self) -> None:
+        with self._state_lock:
+            if self._closed_by_client:
+                return
+            self._closed_by_client = True
+            loop = self._loop
+            run_task = self._run_task
+            websocket = self._websocket
+
+        if loop is not None and websocket is not None and not self._done.is_set():
+            future = asyncio.run_coroutine_threadsafe(websocket.close(), loop)
+            try:
+                future.result(timeout=_WRITE_TIMEOUT)
+            except Exception:
+                pass
+        elif loop is not None and run_task is not None and not run_task.done():
+            loop.call_soon_threadsafe(run_task.cancel)
+        if self._thread.is_alive() and threading.current_thread() is not self._thread:
+            self._thread.join(timeout=_WRITE_TIMEOUT)
+
+    def _send(self, data: bytes | str) -> None:
+        if self._done.is_set():
+            self._raise_error()
+            raise _PtyTransportError("PTY session is no longer running")
+        if threading.current_thread() is self._thread:
+            raise _PtyTransportError(
+                "PTY methods cannot be called from the on_data callback"
+            )
+        with self._state_lock:
+            loop = self._loop
+            websocket = self._websocket
+        if loop is None or websocket is None:
+            raise _PtyTransportError("PTY transport is not connected")
+        future: Future[Any] = asyncio.run_coroutine_threadsafe(
+            websocket.send(data), loop
+        )
+        try:
+            future.result(timeout=_WRITE_TIMEOUT)
+        except Exception as error:
+            raise _PtyTransportError(f"failed to write to PTY: {error}") from error
+
+    def _thread_main(self) -> None:
+        try:
+            asyncio.run(self._run())
+        except BaseException as error:
+            self._set_error(error)
+        finally:
+            self._ready.set()
+            self._done.set()
+            try:
+                self._on_done()
+            except Exception:
+                pass
+
+    async def _run(self) -> None:
+        import websockets
+
+        loop = asyncio.get_running_loop()
+        with self._state_lock:
+            self._loop = loop
+            self._run_task = asyncio.current_task()
+
+        try:
+            async with websockets.connect(
+                self._uri,
+                ssl=self._ssl_context,
+                ping_interval=20,
+                ping_timeout=10,
+            ) as websocket:
+                with self._state_lock:
+                    self._websocket = websocket
+                await websocket.send(f"RESIZE:{self._cols}:{self._rows}")
+
+                async for message in websocket:
+                    if isinstance(message, bytes):
+                        if self._on_data is not None:
+                            try:
+                                self._on_data(message)
+                            except BaseException as error:
+                                self._set_error(
+                                    _PtyTransportError(
+                                        f"PTY on_data callback failed: {error}"
+                                    )
+                                )
+                                await websocket.close()
+                                return
+                        continue
+                    self._handle_control(message)
+                    if self._done.is_set():
+                        await websocket.close()
+                        await websocket.wait_closed()
+                        return
+        except BaseException as error:
+            if not self._closed_by_client and not self._done.is_set():
+                self._set_error(
+                    _PtyTransportError(f"PTY WebSocket connection failed: {error}")
+                )
+        finally:
+            with self._state_lock:
+                self._websocket = None
+                self._loop = None
+                self._run_task = None
+
+        if not self._closed_by_client and not self._done.is_set():
+            self._set_error(
+                _PtyTransportError("PTY connection closed before terminal status")
+            )
+
+    def _handle_control(self, message: str) -> None:
+        try:
+            event = json.loads(message)
+        except json.JSONDecodeError as error:
+            self._set_error(_PtyTransportError(f"invalid PTY control frame: {error}"))
+            return
+        if not isinstance(event, dict) or event.get("version") != _PROTOCOL_VERSION:
+            self._set_error(_PtyTransportError("unsupported PTY control frame"))
+            return
+
+        event_type = event.get("type")
+        session_id = event.get("session_id")
+        if not isinstance(session_id, str) or not session_id:
+            self._set_error(_PtyTransportError("PTY control frame has no session_id"))
+            return
+
+        if event_type == "started":
+            self._session_id = session_id
+            self._ready.set()
+        elif event_type == "exited":
+            exit_code = event.get("exit_code")
+            if isinstance(exit_code, bool) or not isinstance(exit_code, int):
+                self._set_error(_PtyTransportError("PTY exit frame has no exit_code"))
+                return
+            self._session_id = session_id
+            self._exit_code = exit_code
+            self._ready.set()
+            self._done.set()
+        elif event_type == "error":
+            message_value = event.get("message")
+            detail = (
+                message_value if isinstance(message_value, str) else "unknown error"
+            )
+            self._session_id = session_id
+            self._set_error(_PtyTransportError(f"remote PTY error: {detail}"))
+        else:
+            self._set_error(
+                _PtyTransportError(f"unknown PTY control event {event_type!r}")
+            )
+
+    def _set_error(self, error: BaseException) -> None:
+        with self._state_lock:
+            if self._error is None:
+                self._error = error
+        self._ready.set()
+        self._done.set()
+
+    def _raise_error(self) -> None:
+        with self._state_lock:
+            error = self._error
+        if error is not None:
+            raise error