txl-remote-kernels 0.3.7__py3-none-any.whl

txl_remote_kernels/driver.py

@@ -0,0 +1,124 @@
+import json
+import time
+import uuid
+from typing import Any, Dict
+from urllib import parse
+
+import httpx
+from anyio import Lock, sleep
+from anyioutils import create_task
+from httpx_ws import aconnect_ws
+from txl_kernel.driver import KernelMixin
+from txl_kernel.message import date_to_str
+
+from .message import (
+    deserialize_msg_from_ws_v1,
+    from_binary,
+    serialize_msg_to_ws_v1,
+    to_binary,
+)
+
+
+def deadline_to_timeout(deadline: float) -> float:
+    return max(0, deadline - time.time())
+
+
+class KernelDriver(KernelMixin):
+    def __init__(
+        self,
+        task_group,
+        url: str,
+        kernel_name: str | None = "",
+        comm_handlers=[],
+    ) -> None:
+        super().__init__(task_group)
+        self.task_group = task_group
+        self.kernel_name = kernel_name
+        parsed_url = parse.urlparse(url)
+        self.base_url = parse.urljoin(url, parsed_url.path).rstrip("/")
+        self.query_params = parse.parse_qs(parsed_url.query)
+        self.cookies = httpx.Cookies()
+        i = self.base_url.find(":")
+        self.ws_url = ("wss" if self.base_url[i - 1] == "s" else "ws") + self.base_url[i:]
+        self.start_task = create_task(self.start(), task_group)
+        self.comm_handlers = comm_handlers
+        self.shell_channel = "shell"
+        self.control_channel = "control"
+        self.iopub_channel = "iopub"
+        self.send_lock = Lock()
+        self.kernel_id = None
+
+    async def start(self):
+        i = str(uuid.uuid4())
+        async with httpx.AsyncClient() as client:
+            body = {
+                "kernel": {"name": self.kernel_name},
+                "name": i,
+                "path": i,
+                "type": "notebook",
+            }
+            r = await client.post(
+                f"{self.base_url}/api/sessions",
+                json=body,
+                params={**self.query_params},
+                cookies=self.cookies,
+            )
+            d = r.json()
+            self.cookies.update(r.cookies)
+            self.session_id = d["id"]
+            self.kernel_id = d["kernel"]["id"]
+            r = await client.get(
+                f"{self.base_url}/api/kernels/{self.kernel_id}",
+                cookies=self.cookies,
+            )
+        if r.status_code != 200 or self.kernel_id != r.json()["id"]:
+            return
+        async with aconnect_ws(
+            f"{self.ws_url}/api/kernels/{self.kernel_id}/channels",
+            params={"session_id": self.session_id},
+            cookies=self.cookies,
+            subprotocols=["v1.kernel.websocket.jupyter.org"],
+        ) as self.websocket:
+            recv_task = create_task(self._recv(), self.task_group)
+            try:
+                await self.wait_for_ready()
+                self.started.set()
+                await sleep(float("inf"))
+            except BaseException:
+                recv_task.cancel()
+                self.start_task.cancel()
+
+    async def _recv(self):
+        while True:
+            message = await self.websocket.receive()
+            if self.websocket.subprotocol == "v1.kernel.websocket.jupyter.org":
+                msg = deserialize_msg_from_ws_v1(message.data)
+            else:
+                if isinstance(message.data, str):
+                    msg = json.loads(message.data)
+                else:
+                    msg = from_binary(message.data)
+            self.recv_queue.put_nowait(msg)
+
+    async def send_message(
+        self,
+        msg: Dict[str, Any],
+        channel,
+        change_date_to_str: bool = False,
+    ):
+        async with self.send_lock:
+            _date_to_str = date_to_str if change_date_to_str else lambda x: x
+            msg["header"] = _date_to_str(msg["header"])
+            msg["parent_header"] = _date_to_str(msg["parent_header"])
+            msg["metadata"] = _date_to_str(msg["metadata"])
+            msg["content"] = _date_to_str(msg.get("content", {}))
+            msg["channel"] = channel
+            if self.websocket.subprotocol == "v1.kernel.websocket.jupyter.org":
+                bmsg = serialize_msg_to_ws_v1(msg)
+                await self.websocket.send_bytes(bmsg)
+            else:
+                bmsg = to_binary(msg)
+                if bmsg is None:
+                    await self.websocket.send_json(msg)
+                else:
+                    await self.websocket.send_bytes(bmsg)

txl_remote_kernels/main.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+from typing import Any
+from urllib import parse
+
+import httpx
+from anyio import create_task_group, sleep
+from fps import Module
+from pycrdt import Map
+
+from txl.base import Kernels, Kernelspecs
+
+from .driver import KernelDriver
+
+
+class RemoteKernels(Kernels):
+    comm_handlers = []
+
+    def __init__(
+        self,
+        url: str,
+        kernel_name: str | None,
+    ):
+        self.kernel = KernelDriver(
+            self.task_group, url, kernel_name, comm_handlers=self.comm_handlers
+        )
+
+    async def execute(self, ycell: Map):
+        await self.kernel.execute(ycell)
+
+
+class RemoteKernelspecs(Kernelspecs):
+    def __init__(
+        self,
+        url: str,
+    ):
+        parsed_url = parse.urlparse(url)
+        self.base_url = parse.urljoin(url, parsed_url.path).rstrip("/")
+        self.query_params = parse.parse_qs(parsed_url.query)
+        self.cookies = httpx.Cookies()
+
+    async def get(self) -> dict[str, Any]:
+        url = f"{self.base_url}/api/kernelspecs"
+        try:
+            async with httpx.AsyncClient() as client:
+                r = await client.get(
+                    url,
+                    params={**self.query_params},
+                    cookies=self.cookies,
+                )
+                d = r.json()
+                self.cookies.update(r.cookies)
+                return d
+        except httpx.ConnectError:
+            raise RuntimeError(f"Could not connect to a Jupyter server at {url}")
+
+
+class RemoteKernelsModule(Module):
+    def __init__(self, name: str, url: str = "http://127.0.0.1:8000"):
+        super().__init__(name)
+        self.url = url
+
+    async def start(self) -> None:
+        url = self.url
+
+        async with create_task_group() as self.tg:
+            class _RemoteKernels(RemoteKernels):
+                task_group = self.tg
+
+                def __init__(self, *args, **kwargs):
+                    super().__init__(url, *args, **kwargs)
+
+            self.put(_RemoteKernels, Kernels)
+            self.done()
+            await sleep(float("inf"))
+
+    async def stop(self) -> None:
+        self.tg.cancel_scope.cancel()
+
+
+class RemoteKernelspecsModule(Module):
+    def __init__(self, name: str, url: str = "http://127.0.0.1:8000"):
+        super().__init__(name)
+        self.url = url
+
+    async def start(self) -> None:
+        kernelspecs = RemoteKernelspecs(self.url)
+        self.put(kernelspecs, Kernelspecs)

txl_remote_kernels/message.py

@@ -0,0 +1,67 @@
+import json
+import struct
+from typing import Any, Dict, List, Optional
+
+
+def deserialize_msg_from_ws_v1(ws_msg: bytes) -> Dict[str, Any]:
+    offset_number = int.from_bytes(ws_msg[:8], "little")
+    offsets = [
+        int.from_bytes(ws_msg[8 * (i + 1) : 8 * (i + 2)], "little") for i in range(offset_number)
+    ]
+    channel = ws_msg[offsets[0] : offsets[1]].decode("utf-8")
+    msg_list = [ws_msg[offsets[i] : offsets[i + 1]] for i in range(1, offset_number - 1)]
+    msg = {
+        "channel": channel,
+        "header": json.loads(msg_list[0]),
+        "parent_header": json.loads(msg_list[1]),
+        "metadata": json.loads(msg_list[2]),
+        "content": json.loads(msg_list[3]),
+    }
+    msg["buffers"] = msg_list[4:]
+    return msg
+
+
+def serialize_msg_to_ws_v1(msg: Dict[str, Any]) -> List[bytes]:
+    msg_list = [
+        json.dumps(msg["header"]).encode("utf-8"),
+        json.dumps(msg["parent_header"]).encode("utf-8"),
+        json.dumps(msg["metadata"]).encode("utf-8"),
+        json.dumps(msg["content"]).encode("utf-8"),
+    ] + msg.get("buffers", [])
+    channel_b = msg["channel"].encode("utf-8")
+    offsets = []
+    offsets.append(8 * (1 + 1 + len(msg_list) + 1))
+    offsets.append(len(channel_b) + offsets[-1])
+    for msg in msg_list:
+        offsets.append(len(msg) + offsets[-1])
+    offset_number = len(offsets).to_bytes(8, byteorder="little")
+    offsets_b = [offset.to_bytes(8, byteorder="little") for offset in offsets]
+    bin_msg = [offset_number] + offsets_b + [channel_b] + msg_list
+    return b"".join(bin_msg)
+
+
+def to_binary(msg: Dict[str, Any]) -> Optional[bytes]:
+    if not msg["buffers"]:
+        return None
+    buffers = msg.pop("buffers")
+    bmsg = json.dumps(msg).encode("utf8")
+    buffers.insert(0, bmsg)
+    n = len(buffers)
+    offsets = [4 * (n + 1)]
+    for b in buffers[:-1]:
+        offsets.append(offsets[-1] + len(b))
+    header = struct.pack("!" + "I" * (n + 1), n, *offsets)
+    buffers.insert(0, header)
+    return b"".join(buffers)
+
+
+def from_binary(bmsg: bytes) -> Dict[str, Any]:
+    n = struct.unpack("!i", bmsg[:4])[0]
+    offsets = list(struct.unpack("!" + "I" * n, bmsg[4 : 4 * (n + 1)]))  # noqa
+    offsets.append(None)
+    buffers = []
+    for start, stop in zip(offsets[:-1], offsets[1:]):
+        buffers.append(bmsg[start:stop])
+    msg = json.loads(buffers[0].decode("utf8"))
+    msg["buffers"] = buffers[1:]
+    return msg

txl_remote_kernels-0.3.7.dist-info/METADATA

@@ -0,0 +1,25 @@
+Metadata-Version: 2.4
+Name: txl_remote_kernels
+Version: 0.3.7
+Summary: TXL plugin for remote kernels
+Project-URL: Source, https://github.com/davidbrochart/jpterm/tree/main/plugins/remote_kernels
+Author-email: David Brochart <david.brochart@gmail.com>
+License-Expression: MIT
+License-File: LICENSE.txt
+Classifier: Development Status :: 4 - Beta
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: Programming Language :: Python :: Implementation :: PyPy
+Requires-Python: >=3.10
+Requires-Dist: httpx-ws>=0.4.2
+Requires-Dist: httpx>=0.23.1
+Requires-Dist: pycrdt<0.13.0,>=0.12.44
+Requires-Dist: txl-kernel
+Requires-Dist: txl==0.3.3
+Description-Content-Type: text/markdown
+
+# TXL plugin for remote kernels

txl_remote_kernels-0.3.7.dist-info/RECORD

@@ -0,0 +1,9 @@
+txl_remote_kernels/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+txl_remote_kernels/driver.py,sha256=gJuGmbshysBo-040NThjm-MamqIME47ejUoeQuoGMmY,4332
+txl_remote_kernels/main.py,sha256=L7oQW6d0WKu35-luLrZwAWoml4X7CH0OjFf6ASrA9nE,2429
+txl_remote_kernels/message.py,sha256=lyO6uPpRlYgdRSOZxam6iA-QH_HjXga3XKNlA_cCs5Q,2400
+txl_remote_kernels-0.3.7.dist-info/METADATA,sha256=NX5tMllWLtay7_AxzMtOkXx_42JxG3h3aW5OEs1uAII,980
+txl_remote_kernels-0.3.7.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+txl_remote_kernels-0.3.7.dist-info/entry_points.txt,sha256=WWe5nk2g-g6A4T9-ceUej1JR4k4rh4KF4FIbriEFQVM,289
+txl_remote_kernels-0.3.7.dist-info/licenses/LICENSE.txt,sha256=su0IgzSHZ9tMFeCQ_IaCji5Q4VrZSOqDsEFi7lPduc8,1106
+txl_remote_kernels-0.3.7.dist-info/RECORD,,

txl_remote_kernels-0.3.7.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any

txl_remote_kernels-0.3.7.dist-info/entry_points.txt

@@ -0,0 +1,7 @@
+[fps.modules]
+remote_kernels = txl_remote_kernels.main:RemoteKernelsModule
+remote_kernelspecs = txl_remote_kernels.main:RemoteKernelspecsModule
+
+[txl.modules]
+remote_kernels = txl_remote_kernels.main:RemoteKernelsModule
+remote_kernelspecs = txl_remote_kernels.main:RemoteKernelspecsModule

txl_remote_kernels-0.3.7.dist-info/licenses/LICENSE.txt

@@ -0,0 +1,9 @@
+MIT License
+
+Copyright (c) 2022-present David Brochart <david.brochart@gmail.com>
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.