txl_remote_kernels 0.1.5__tar.gz → 0.2.4__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- txl_remote_kernels-0.2.4/.gitignore +3 -0
- {txl_remote_kernels-0.1.5 → txl_remote_kernels-0.2.4}/PKG-INFO +6 -8
- {txl_remote_kernels-0.1.5 → txl_remote_kernels-0.2.4}/pyproject.toml +8 -9
- txl_remote_kernels-0.2.4/txl_remote_kernels/__init__.py +1 -0
- txl_remote_kernels-0.2.4/txl_remote_kernels/components.py +83 -0
- txl_remote_kernels-0.2.4/txl_remote_kernels/driver.py +121 -0
- txl_remote_kernels-0.2.4/txl_remote_kernels/message.py +67 -0
- txl_remote_kernels-0.1.5/.gitignore +0 -1
- txl_remote_kernels-0.1.5/txl_remote_kernels/__init__.py +0 -1
- txl_remote_kernels-0.1.5/txl_remote_kernels/components.py +0 -37
- txl_remote_kernels-0.1.5/txl_remote_kernels/driver.py +0 -237
- txl_remote_kernels-0.1.5/txl_remote_kernels/message.py +0 -108
- {txl_remote_kernels-0.1.5 → txl_remote_kernels-0.2.4}/LICENSE.txt +0 -0
- {txl_remote_kernels-0.1.5 → txl_remote_kernels-0.2.4}/README.md +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.1
|
|
2
2
|
Name: txl_remote_kernels
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.2.4
|
|
4
4
|
Summary: TXL plugin for remote kernels
|
|
5
5
|
Project-URL: Source, https://github.com/davidbrochart/jpterm/plugins/remote_kernels
|
|
6
6
|
Author-email: David Brochart <david.brochart@gmail.com>
|
|
@@ -8,19 +8,17 @@ License-Expression: MIT
|
|
|
8
8
|
License-File: LICENSE.txt
|
|
9
9
|
Classifier: Development Status :: 4 - Beta
|
|
10
10
|
Classifier: Programming Language :: Python
|
|
11
|
-
Classifier: Programming Language :: Python :: 3.7
|
|
12
|
-
Classifier: Programming Language :: Python :: 3.8
|
|
13
|
-
Classifier: Programming Language :: Python :: 3.9
|
|
14
11
|
Classifier: Programming Language :: Python :: 3.10
|
|
15
12
|
Classifier: Programming Language :: Python :: 3.11
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
16
14
|
Classifier: Programming Language :: Python :: Implementation :: CPython
|
|
17
15
|
Classifier: Programming Language :: Python :: Implementation :: PyPy
|
|
18
|
-
Requires-Python: >=3.
|
|
19
|
-
Requires-Dist: httpx-ws>=0.2
|
|
16
|
+
Requires-Python: >=3.10
|
|
17
|
+
Requires-Dist: httpx-ws>=0.4.2
|
|
20
18
|
Requires-Dist: httpx>=0.23.1
|
|
21
|
-
Requires-Dist:
|
|
19
|
+
Requires-Dist: pycrdt<0.8.0,>=0.7.2
|
|
22
20
|
Requires-Dist: txl
|
|
23
|
-
Requires-Dist:
|
|
21
|
+
Requires-Dist: txl-kernel
|
|
24
22
|
Description-Content-Type: text/markdown
|
|
25
23
|
|
|
26
24
|
# TXL plugin for remote kernels
|
|
@@ -6,7 +6,7 @@ build-backend = "hatchling.build"
|
|
|
6
6
|
name = "txl_remote_kernels"
|
|
7
7
|
description = "TXL plugin for remote kernels"
|
|
8
8
|
readme = "README.md"
|
|
9
|
-
requires-python = ">=3.
|
|
9
|
+
requires-python = ">=3.10"
|
|
10
10
|
license = "MIT"
|
|
11
11
|
keywords = []
|
|
12
12
|
authors = [
|
|
@@ -15,28 +15,27 @@ authors = [
|
|
|
15
15
|
classifiers = [
|
|
16
16
|
"Development Status :: 4 - Beta",
|
|
17
17
|
"Programming Language :: Python",
|
|
18
|
-
"Programming Language :: Python :: 3.7",
|
|
19
|
-
"Programming Language :: Python :: 3.8",
|
|
20
|
-
"Programming Language :: Python :: 3.9",
|
|
21
18
|
"Programming Language :: Python :: 3.10",
|
|
22
19
|
"Programming Language :: Python :: 3.11",
|
|
20
|
+
"Programming Language :: Python :: 3.12",
|
|
23
21
|
"Programming Language :: Python :: Implementation :: CPython",
|
|
24
22
|
"Programming Language :: Python :: Implementation :: PyPy",
|
|
25
23
|
]
|
|
26
24
|
dependencies = [
|
|
27
25
|
"txl",
|
|
28
|
-
"
|
|
26
|
+
"txl_kernel",
|
|
29
27
|
"httpx >=0.23.1",
|
|
30
|
-
"httpx-ws >=0.2
|
|
31
|
-
"
|
|
28
|
+
"httpx-ws >=0.4.2",
|
|
29
|
+
"pycrdt >=0.7.2,<0.8.0",
|
|
32
30
|
]
|
|
33
31
|
dynamic = ["version"]
|
|
34
32
|
|
|
35
33
|
[project.urls]
|
|
36
34
|
Source = "https://github.com/davidbrochart/jpterm/plugins/remote_kernels"
|
|
37
35
|
|
|
38
|
-
[project.entry-points
|
|
39
|
-
txl_remote_kernels = "txl_remote_kernels.components"
|
|
36
|
+
[project.entry-points]
|
|
37
|
+
"asphalt.components" = {remote_kernels = "txl_remote_kernels.components:RemoteKernelsComponent", remote_kernelspecs = "txl_remote_kernels.components:RemoteKernelspecsComponent"}
|
|
38
|
+
"txl.components" = {remote_kernels = "txl_remote_kernels.components:RemoteKernelsComponent", remote_kernelspecs = "txl_remote_kernels.components:RemoteKernelspecsComponent"}
|
|
40
39
|
|
|
41
40
|
[tool.hatch.version]
|
|
42
41
|
path = "txl_remote_kernels/__init__.py"
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
__version__ = "0.2.4"
|
|
@@ -0,0 +1,83 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from typing import Any
|
|
4
|
+
from urllib import parse
|
|
5
|
+
|
|
6
|
+
import httpx
|
|
7
|
+
from asphalt.core import Component, Context
|
|
8
|
+
from pycrdt import Map
|
|
9
|
+
|
|
10
|
+
from txl.base import Kernels, Kernelspecs
|
|
11
|
+
|
|
12
|
+
from .driver import KernelDriver
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class RemoteKernels(Kernels):
|
|
16
|
+
comm_handlers = []
|
|
17
|
+
|
|
18
|
+
def __init__(
|
|
19
|
+
self,
|
|
20
|
+
url: str,
|
|
21
|
+
kernel_name: str | None,
|
|
22
|
+
):
|
|
23
|
+
self.kernel = KernelDriver(url, kernel_name, comm_handlers=self.comm_handlers)
|
|
24
|
+
|
|
25
|
+
async def execute(self, ycell: Map):
|
|
26
|
+
await self.kernel.execute(ycell)
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
class RemoteKernelspecs(Kernelspecs):
|
|
30
|
+
def __init__(
|
|
31
|
+
self,
|
|
32
|
+
url: str,
|
|
33
|
+
):
|
|
34
|
+
parsed_url = parse.urlparse(url)
|
|
35
|
+
self.base_url = parse.urljoin(url, parsed_url.path).rstrip("/")
|
|
36
|
+
self.query_params = parse.parse_qs(parsed_url.query)
|
|
37
|
+
self.cookies = httpx.Cookies()
|
|
38
|
+
|
|
39
|
+
async def get(self) -> dict[str, Any]:
|
|
40
|
+
url = f"{self.base_url}/api/kernelspecs"
|
|
41
|
+
try:
|
|
42
|
+
async with httpx.AsyncClient() as client:
|
|
43
|
+
r = await client.get(
|
|
44
|
+
url,
|
|
45
|
+
params={**self.query_params},
|
|
46
|
+
cookies=self.cookies,
|
|
47
|
+
)
|
|
48
|
+
d = r.json()
|
|
49
|
+
self.cookies.update(r.cookies)
|
|
50
|
+
return d
|
|
51
|
+
except httpx.ConnectError:
|
|
52
|
+
raise RuntimeError(f"Could not connect to a Jupyter server at {url}")
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
class RemoteKernelsComponent(Component):
|
|
56
|
+
def __init__(self, url: str = "http://127.0.0.1:8000"):
|
|
57
|
+
super().__init__()
|
|
58
|
+
self.url = url
|
|
59
|
+
|
|
60
|
+
async def start(
|
|
61
|
+
self,
|
|
62
|
+
ctx: Context,
|
|
63
|
+
) -> None:
|
|
64
|
+
url = self.url
|
|
65
|
+
|
|
66
|
+
class _RemoteKernels(RemoteKernels):
|
|
67
|
+
def __init__(self, *args, **kwargs):
|
|
68
|
+
super().__init__(url, *args, **kwargs)
|
|
69
|
+
|
|
70
|
+
ctx.add_resource(_RemoteKernels, types=Kernels)
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
class RemoteKernelspecsComponent(Component):
|
|
74
|
+
def __init__(self, url: str = "http://127.0.0.1:8000"):
|
|
75
|
+
super().__init__()
|
|
76
|
+
self.url = url
|
|
77
|
+
|
|
78
|
+
async def start(
|
|
79
|
+
self,
|
|
80
|
+
ctx: Context,
|
|
81
|
+
) -> None:
|
|
82
|
+
kernelspecs = RemoteKernelspecs(self.url)
|
|
83
|
+
ctx.add_resource(kernelspecs, types=Kernelspecs)
|
|
@@ -0,0 +1,121 @@
|
|
|
1
|
+
import asyncio
|
|
2
|
+
import json
|
|
3
|
+
import time
|
|
4
|
+
import uuid
|
|
5
|
+
from typing import Any, Dict
|
|
6
|
+
from urllib import parse
|
|
7
|
+
|
|
8
|
+
import httpx
|
|
9
|
+
from httpx_ws import aconnect_ws
|
|
10
|
+
from txl_kernel.driver import KernelMixin
|
|
11
|
+
from txl_kernel.message import date_to_str
|
|
12
|
+
|
|
13
|
+
from .message import (
|
|
14
|
+
deserialize_msg_from_ws_v1,
|
|
15
|
+
from_binary,
|
|
16
|
+
serialize_msg_to_ws_v1,
|
|
17
|
+
to_binary,
|
|
18
|
+
)
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def deadline_to_timeout(deadline: float) -> float:
|
|
22
|
+
return max(0, deadline - time.time())
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class KernelDriver(KernelMixin):
|
|
26
|
+
def __init__(
|
|
27
|
+
self,
|
|
28
|
+
url: str,
|
|
29
|
+
kernel_name: str | None = "",
|
|
30
|
+
comm_handlers=[],
|
|
31
|
+
) -> None:
|
|
32
|
+
super().__init__()
|
|
33
|
+
self.kernel_name = kernel_name
|
|
34
|
+
parsed_url = parse.urlparse(url)
|
|
35
|
+
self.base_url = parse.urljoin(url, parsed_url.path).rstrip("/")
|
|
36
|
+
self.query_params = parse.parse_qs(parsed_url.query)
|
|
37
|
+
self.cookies = httpx.Cookies()
|
|
38
|
+
i = self.base_url.find(":")
|
|
39
|
+
self.ws_url = ("wss" if self.base_url[i - 1] == "s" else "ws") + self.base_url[i:]
|
|
40
|
+
self.start_task = asyncio.create_task(self.start())
|
|
41
|
+
self.comm_handlers = comm_handlers
|
|
42
|
+
self.shell_channel = "shell"
|
|
43
|
+
self.control_channel = "control"
|
|
44
|
+
self.iopub_channel = "iopub"
|
|
45
|
+
self.send_lock = asyncio.Lock()
|
|
46
|
+
self.kernel_id = None
|
|
47
|
+
|
|
48
|
+
async def start(self):
|
|
49
|
+
i = str(uuid.uuid4())
|
|
50
|
+
async with httpx.AsyncClient() as client:
|
|
51
|
+
body = {
|
|
52
|
+
"kernel": {"name": self.kernel_name},
|
|
53
|
+
"name": i,
|
|
54
|
+
"path": i,
|
|
55
|
+
"type": "notebook",
|
|
56
|
+
}
|
|
57
|
+
r = await client.post(
|
|
58
|
+
f"{self.base_url}/api/sessions",
|
|
59
|
+
json=body,
|
|
60
|
+
params={**self.query_params},
|
|
61
|
+
cookies=self.cookies,
|
|
62
|
+
)
|
|
63
|
+
d = r.json()
|
|
64
|
+
self.cookies.update(r.cookies)
|
|
65
|
+
self.session_id = d["id"]
|
|
66
|
+
self.kernel_id = d["kernel"]["id"]
|
|
67
|
+
r = await client.get(
|
|
68
|
+
f"{self.base_url}/api/kernels/{self.kernel_id}",
|
|
69
|
+
cookies=self.cookies,
|
|
70
|
+
)
|
|
71
|
+
if r.status_code != 200 or self.kernel_id != r.json()["id"]:
|
|
72
|
+
return
|
|
73
|
+
async with aconnect_ws(
|
|
74
|
+
f"{self.ws_url}/api/kernels/{self.kernel_id}/channels",
|
|
75
|
+
params={"session_id": self.session_id},
|
|
76
|
+
cookies=self.cookies,
|
|
77
|
+
subprotocols=["v1.kernel.websocket.jupyter.org"],
|
|
78
|
+
) as self.websocket:
|
|
79
|
+
recv_task = asyncio.create_task(self._recv())
|
|
80
|
+
try:
|
|
81
|
+
await self.wait_for_ready()
|
|
82
|
+
self.started.set()
|
|
83
|
+
await asyncio.Future()
|
|
84
|
+
except BaseException:
|
|
85
|
+
recv_task.cancel()
|
|
86
|
+
self.start_task.cancel()
|
|
87
|
+
|
|
88
|
+
async def _recv(self):
|
|
89
|
+
while True:
|
|
90
|
+
message = await self.websocket.receive()
|
|
91
|
+
if self.websocket.subprotocol == "v1.kernel.websocket.jupyter.org":
|
|
92
|
+
msg = deserialize_msg_from_ws_v1(message.data)
|
|
93
|
+
else:
|
|
94
|
+
if isinstance(message.data, str):
|
|
95
|
+
msg = json.loads(message.data)
|
|
96
|
+
else:
|
|
97
|
+
msg = from_binary(message.data)
|
|
98
|
+
self.recv_queue.put_nowait(msg)
|
|
99
|
+
|
|
100
|
+
async def send_message(
|
|
101
|
+
self,
|
|
102
|
+
msg: Dict[str, Any],
|
|
103
|
+
channel,
|
|
104
|
+
change_date_to_str: bool = False,
|
|
105
|
+
):
|
|
106
|
+
async with self.send_lock:
|
|
107
|
+
_date_to_str = date_to_str if change_date_to_str else lambda x: x
|
|
108
|
+
msg["header"] = _date_to_str(msg["header"])
|
|
109
|
+
msg["parent_header"] = _date_to_str(msg["parent_header"])
|
|
110
|
+
msg["metadata"] = _date_to_str(msg["metadata"])
|
|
111
|
+
msg["content"] = _date_to_str(msg.get("content", {}))
|
|
112
|
+
msg["channel"] = channel
|
|
113
|
+
if self.websocket.subprotocol == "v1.kernel.websocket.jupyter.org":
|
|
114
|
+
bmsg = serialize_msg_to_ws_v1(msg)
|
|
115
|
+
await self.websocket.send_bytes(bmsg)
|
|
116
|
+
else:
|
|
117
|
+
bmsg = to_binary(msg)
|
|
118
|
+
if bmsg is None:
|
|
119
|
+
await self.websocket.send_json(msg)
|
|
120
|
+
else:
|
|
121
|
+
await self.websocket.send_bytes(bmsg)
|
|
@@ -0,0 +1,67 @@
|
|
|
1
|
+
import json
|
|
2
|
+
import struct
|
|
3
|
+
from typing import Any, Dict, List, Optional
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
def deserialize_msg_from_ws_v1(ws_msg: bytes) -> Dict[str, Any]:
|
|
7
|
+
offset_number = int.from_bytes(ws_msg[:8], "little")
|
|
8
|
+
offsets = [
|
|
9
|
+
int.from_bytes(ws_msg[8 * (i + 1) : 8 * (i + 2)], "little") for i in range(offset_number)
|
|
10
|
+
]
|
|
11
|
+
channel = ws_msg[offsets[0] : offsets[1]].decode("utf-8")
|
|
12
|
+
msg_list = [ws_msg[offsets[i] : offsets[i + 1]] for i in range(1, offset_number - 1)]
|
|
13
|
+
msg = {
|
|
14
|
+
"channel": channel,
|
|
15
|
+
"header": json.loads(msg_list[0]),
|
|
16
|
+
"parent_header": json.loads(msg_list[1]),
|
|
17
|
+
"metadata": json.loads(msg_list[2]),
|
|
18
|
+
"content": json.loads(msg_list[3]),
|
|
19
|
+
}
|
|
20
|
+
msg["buffers"] = msg_list[4:]
|
|
21
|
+
return msg
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def serialize_msg_to_ws_v1(msg: Dict[str, Any]) -> List[bytes]:
|
|
25
|
+
msg_list = [
|
|
26
|
+
json.dumps(msg["header"]).encode("utf-8"),
|
|
27
|
+
json.dumps(msg["parent_header"]).encode("utf-8"),
|
|
28
|
+
json.dumps(msg["metadata"]).encode("utf-8"),
|
|
29
|
+
json.dumps(msg["content"]).encode("utf-8"),
|
|
30
|
+
] + msg.get("buffers", [])
|
|
31
|
+
channel_b = msg["channel"].encode("utf-8")
|
|
32
|
+
offsets = []
|
|
33
|
+
offsets.append(8 * (1 + 1 + len(msg_list) + 1))
|
|
34
|
+
offsets.append(len(channel_b) + offsets[-1])
|
|
35
|
+
for msg in msg_list:
|
|
36
|
+
offsets.append(len(msg) + offsets[-1])
|
|
37
|
+
offset_number = len(offsets).to_bytes(8, byteorder="little")
|
|
38
|
+
offsets_b = [offset.to_bytes(8, byteorder="little") for offset in offsets]
|
|
39
|
+
bin_msg = [offset_number] + offsets_b + [channel_b] + msg_list
|
|
40
|
+
return b"".join(bin_msg)
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
def to_binary(msg: Dict[str, Any]) -> Optional[bytes]:
|
|
44
|
+
if not msg["buffers"]:
|
|
45
|
+
return None
|
|
46
|
+
buffers = msg.pop("buffers")
|
|
47
|
+
bmsg = json.dumps(msg).encode("utf8")
|
|
48
|
+
buffers.insert(0, bmsg)
|
|
49
|
+
n = len(buffers)
|
|
50
|
+
offsets = [4 * (n + 1)]
|
|
51
|
+
for b in buffers[:-1]:
|
|
52
|
+
offsets.append(offsets[-1] + len(b))
|
|
53
|
+
header = struct.pack("!" + "I" * (n + 1), n, *offsets)
|
|
54
|
+
buffers.insert(0, header)
|
|
55
|
+
return b"".join(buffers)
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def from_binary(bmsg: bytes) -> Dict[str, Any]:
|
|
59
|
+
n = struct.unpack("!i", bmsg[:4])[0]
|
|
60
|
+
offsets = list(struct.unpack("!" + "I" * n, bmsg[4 : 4 * (n + 1)])) # noqa
|
|
61
|
+
offsets.append(None)
|
|
62
|
+
buffers = []
|
|
63
|
+
for start, stop in zip(offsets[:-1], offsets[1:]):
|
|
64
|
+
buffers.append(bmsg[start:stop])
|
|
65
|
+
msg = json.loads(buffers[0].decode("utf8"))
|
|
66
|
+
msg["buffers"] = buffers[1:]
|
|
67
|
+
return msg
|
|
@@ -1 +0,0 @@
|
|
|
1
|
-
__pycache__
|
|
@@ -1 +0,0 @@
|
|
|
1
|
-
__version__ = "0.1.5"
|
|
@@ -1,37 +0,0 @@
|
|
|
1
|
-
from functools import partial
|
|
2
|
-
|
|
3
|
-
import y_py as Y
|
|
4
|
-
from asphalt.core import Component, Context
|
|
5
|
-
|
|
6
|
-
from txl.base import Kernels
|
|
7
|
-
from txl.hooks import register_component
|
|
8
|
-
|
|
9
|
-
from .driver import KernelDriver
|
|
10
|
-
|
|
11
|
-
|
|
12
|
-
class RemoteKernels(Kernels):
|
|
13
|
-
def __init__(
|
|
14
|
-
self,
|
|
15
|
-
url: str,
|
|
16
|
-
kernel_name: str,
|
|
17
|
-
):
|
|
18
|
-
self.kernel = KernelDriver(url, kernel_name)
|
|
19
|
-
|
|
20
|
-
async def execute(self, ydoc: Y.YDoc, ycell: Y.YMap):
|
|
21
|
-
await self.kernel.execute(ydoc, ycell)
|
|
22
|
-
|
|
23
|
-
|
|
24
|
-
class RemoteKernelsComponent(Component):
|
|
25
|
-
def __init__(self, url: str = "http://127.0.0.1:8000"):
|
|
26
|
-
super().__init__()
|
|
27
|
-
self.url = url
|
|
28
|
-
|
|
29
|
-
async def start(
|
|
30
|
-
self,
|
|
31
|
-
ctx: Context,
|
|
32
|
-
) -> None:
|
|
33
|
-
terminals_factory = partial(RemoteKernels, self.url)
|
|
34
|
-
ctx.add_resource(terminals_factory, name="kernels", types=Kernels)
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
c = register_component("kernels", RemoteKernelsComponent)
|
|
@@ -1,237 +0,0 @@
|
|
|
1
|
-
import asyncio
|
|
2
|
-
import time
|
|
3
|
-
import uuid
|
|
4
|
-
from typing import Dict, List
|
|
5
|
-
from urllib import parse
|
|
6
|
-
|
|
7
|
-
import httpx
|
|
8
|
-
import y_py as Y
|
|
9
|
-
from httpx_ws import WebSocketNetworkError, aconnect_ws
|
|
10
|
-
|
|
11
|
-
from .message import create_message, send_message, str_to_date
|
|
12
|
-
|
|
13
|
-
|
|
14
|
-
def deadline_to_timeout(deadline: float) -> float:
|
|
15
|
-
return max(0, deadline - time.time())
|
|
16
|
-
|
|
17
|
-
|
|
18
|
-
class KernelDriver:
|
|
19
|
-
def __init__(
|
|
20
|
-
self,
|
|
21
|
-
url: str,
|
|
22
|
-
kernel_name: str = "",
|
|
23
|
-
) -> None:
|
|
24
|
-
self.kernel_name = kernel_name
|
|
25
|
-
parsed_url = parse.urlparse(url)
|
|
26
|
-
self.base_url = parse.urljoin(url, parsed_url.path).rstrip("/")
|
|
27
|
-
self.query_params = parse.parse_qs(parsed_url.query)
|
|
28
|
-
self.cookies = httpx.Cookies()
|
|
29
|
-
i = self.base_url.find(":")
|
|
30
|
-
self.ws_url = ("wss" if self.base_url[i - 1] == "s" else "ws") + self.base_url[
|
|
31
|
-
i:
|
|
32
|
-
]
|
|
33
|
-
self.msg_cnt = 0
|
|
34
|
-
self.execute_requests: Dict[str, Dict[str, List[asyncio.Future]]] = {}
|
|
35
|
-
self.started = asyncio.Event()
|
|
36
|
-
self.start_task = asyncio.create_task(self.start())
|
|
37
|
-
|
|
38
|
-
async def start(self):
|
|
39
|
-
i = str(uuid.uuid4())
|
|
40
|
-
async with httpx.AsyncClient() as client:
|
|
41
|
-
body = {
|
|
42
|
-
"kernel": {"name": self.kernel_name},
|
|
43
|
-
"name": i,
|
|
44
|
-
"path": i,
|
|
45
|
-
"type": "notebook",
|
|
46
|
-
}
|
|
47
|
-
r = await client.post(
|
|
48
|
-
f"{self.base_url}/api/sessions",
|
|
49
|
-
json=body,
|
|
50
|
-
params={**self.query_params},
|
|
51
|
-
cookies=self.cookies,
|
|
52
|
-
)
|
|
53
|
-
response = r.json()
|
|
54
|
-
self.cookies.update(r.cookies)
|
|
55
|
-
self.session_id = response["id"]
|
|
56
|
-
kernel_id = response["kernel"]["id"]
|
|
57
|
-
response = await client.get(
|
|
58
|
-
f"{self.base_url}/api/kernels/{kernel_id}",
|
|
59
|
-
cookies=self.cookies,
|
|
60
|
-
)
|
|
61
|
-
if response.status_code != 200 or kernel_id != response.json()["id"]:
|
|
62
|
-
return
|
|
63
|
-
async with aconnect_ws(
|
|
64
|
-
f"{self.ws_url}/api/kernels/{kernel_id}/channels",
|
|
65
|
-
params={"session_id": self.session_id},
|
|
66
|
-
cookies=self.cookies,
|
|
67
|
-
) as self.websocket:
|
|
68
|
-
recv_task = asyncio.create_task(self._recv())
|
|
69
|
-
try:
|
|
70
|
-
await self._wait_for_ready()
|
|
71
|
-
await asyncio.Future()
|
|
72
|
-
except BaseException:
|
|
73
|
-
recv_task.cancel()
|
|
74
|
-
self.start_task.cancel()
|
|
75
|
-
|
|
76
|
-
async def _recv(self):
|
|
77
|
-
try:
|
|
78
|
-
while True:
|
|
79
|
-
message = await self.websocket.receive_json()
|
|
80
|
-
channel = message.pop("channel")
|
|
81
|
-
message["header"] = str_to_date(message["header"])
|
|
82
|
-
message["parent_header"] = str_to_date(message["parent_header"])
|
|
83
|
-
msg_id = message["parent_header"].get("msg_id")
|
|
84
|
-
if msg_id in self.execute_requests:
|
|
85
|
-
future_messages = self.execute_requests[msg_id][channel]
|
|
86
|
-
if future_messages:
|
|
87
|
-
fut = future_messages[-1]
|
|
88
|
-
if fut.done():
|
|
89
|
-
fut = asyncio.Future()
|
|
90
|
-
future_messages.append(fut)
|
|
91
|
-
else:
|
|
92
|
-
fut = asyncio.Future()
|
|
93
|
-
future_messages.append(fut)
|
|
94
|
-
fut.set_result(message)
|
|
95
|
-
except WebSocketNetworkError:
|
|
96
|
-
pass
|
|
97
|
-
|
|
98
|
-
async def _wait_for_ready(self, timeout: float = float("inf")):
|
|
99
|
-
deadline = time.time() + timeout
|
|
100
|
-
new_timeout = timeout
|
|
101
|
-
while True:
|
|
102
|
-
msg = create_message(
|
|
103
|
-
"kernel_info_request",
|
|
104
|
-
session_id=self.session_id,
|
|
105
|
-
msg_id=str(self.msg_cnt),
|
|
106
|
-
)
|
|
107
|
-
self.msg_cnt += 1
|
|
108
|
-
await send_message(msg, "shell", self.websocket, change_date_to_str=True)
|
|
109
|
-
msg_id = msg["header"]["msg_id"]
|
|
110
|
-
self.execute_requests[msg_id] = {
|
|
111
|
-
"iopub": [asyncio.Future()],
|
|
112
|
-
"shell": [asyncio.Future()],
|
|
113
|
-
}
|
|
114
|
-
try:
|
|
115
|
-
msg = await asyncio.wait_for(
|
|
116
|
-
self.execute_requests[msg_id]["shell"][0], new_timeout
|
|
117
|
-
)
|
|
118
|
-
except asyncio.TimeoutError:
|
|
119
|
-
del self.execute_requests[msg_id]
|
|
120
|
-
error_message = f"Kernel didn't respond in {timeout} seconds"
|
|
121
|
-
raise RuntimeError(error_message)
|
|
122
|
-
if msg["header"]["msg_type"] == "kernel_info_reply":
|
|
123
|
-
try:
|
|
124
|
-
msg = await asyncio.wait_for(
|
|
125
|
-
self.execute_requests[msg_id]["iopub"][0], 0.2
|
|
126
|
-
)
|
|
127
|
-
except asyncio.TimeoutError:
|
|
128
|
-
pass
|
|
129
|
-
else:
|
|
130
|
-
break
|
|
131
|
-
del self.execute_requests[msg_id]
|
|
132
|
-
new_timeout = deadline_to_timeout(deadline)
|
|
133
|
-
self.started.set()
|
|
134
|
-
|
|
135
|
-
async def execute(
|
|
136
|
-
self,
|
|
137
|
-
ydoc: Y.YDoc,
|
|
138
|
-
ycell: Y.YMap,
|
|
139
|
-
timeout: float = float("inf"),
|
|
140
|
-
msg_id: str = "",
|
|
141
|
-
wait_for_executed: bool = True,
|
|
142
|
-
) -> None:
|
|
143
|
-
await self.started.wait()
|
|
144
|
-
if ycell["cell_type"] != "code":
|
|
145
|
-
return
|
|
146
|
-
code = str(ycell["source"])
|
|
147
|
-
content = {"code": code, "silent": False}
|
|
148
|
-
msg = create_message(
|
|
149
|
-
"execute_request",
|
|
150
|
-
content,
|
|
151
|
-
session_id=self.session_id,
|
|
152
|
-
msg_id=str(self.msg_cnt),
|
|
153
|
-
)
|
|
154
|
-
if msg_id:
|
|
155
|
-
msg["header"]["msg_id"] = msg_id
|
|
156
|
-
else:
|
|
157
|
-
msg_id = msg["header"]["msg_id"]
|
|
158
|
-
self.msg_cnt += 1
|
|
159
|
-
await send_message(msg, "shell", self.websocket, change_date_to_str=True)
|
|
160
|
-
if wait_for_executed:
|
|
161
|
-
deadline = time.time() + timeout
|
|
162
|
-
self.execute_requests[msg_id] = {
|
|
163
|
-
"iopub": [asyncio.Future()],
|
|
164
|
-
"shell": [asyncio.Future()],
|
|
165
|
-
}
|
|
166
|
-
try:
|
|
167
|
-
await asyncio.wait_for(
|
|
168
|
-
self.execute_requests[msg_id]["iopub"][0],
|
|
169
|
-
deadline_to_timeout(deadline),
|
|
170
|
-
)
|
|
171
|
-
except asyncio.TimeoutError:
|
|
172
|
-
error_message = f"Kernel didn't respond in {timeout} seconds"
|
|
173
|
-
del self.execute_requests[msg_id]
|
|
174
|
-
raise RuntimeError(error_message)
|
|
175
|
-
await self._handle_outputs(
|
|
176
|
-
ydoc, ycell, self.execute_requests[msg_id]["iopub"]
|
|
177
|
-
)
|
|
178
|
-
try:
|
|
179
|
-
await asyncio.wait_for(
|
|
180
|
-
self.execute_requests[msg_id]["shell"][0],
|
|
181
|
-
deadline_to_timeout(deadline),
|
|
182
|
-
)
|
|
183
|
-
except asyncio.TimeoutError:
|
|
184
|
-
del self.execute_requests[msg_id]
|
|
185
|
-
error_message = f"Kernel didn't respond in {timeout} seconds"
|
|
186
|
-
raise RuntimeError(error_message)
|
|
187
|
-
msg = self.execute_requests[msg_id]["shell"][0].result()
|
|
188
|
-
with ydoc.begin_transaction() as txn:
|
|
189
|
-
ycell.set(txn, "execution_count", msg["content"]["execution_count"])
|
|
190
|
-
del self.execute_requests[msg_id]
|
|
191
|
-
|
|
192
|
-
async def _handle_outputs(
|
|
193
|
-
self, ydoc: Y.YDoc, ycell: Y.YMap, future_messages: List[asyncio.Future]
|
|
194
|
-
):
|
|
195
|
-
with ydoc.begin_transaction() as txn:
|
|
196
|
-
ycell.set(txn, "outputs", [])
|
|
197
|
-
|
|
198
|
-
while True:
|
|
199
|
-
if not future_messages:
|
|
200
|
-
future_messages.append(asyncio.Future())
|
|
201
|
-
fut = future_messages[0]
|
|
202
|
-
if not fut.done():
|
|
203
|
-
await fut
|
|
204
|
-
future_messages.pop(0)
|
|
205
|
-
msg = fut.result()
|
|
206
|
-
msg_type = msg["header"]["msg_type"]
|
|
207
|
-
content = msg["content"]
|
|
208
|
-
outputs = list(ycell["outputs"])
|
|
209
|
-
if msg_type == "stream":
|
|
210
|
-
if (len(outputs) == 0) or (outputs[-1]["name"] != content["name"]):
|
|
211
|
-
outputs.append(
|
|
212
|
-
{"name": content["name"], "output_type": msg_type, "text": []}
|
|
213
|
-
)
|
|
214
|
-
outputs[-1]["text"].append(content["text"])
|
|
215
|
-
elif msg_type in ("display_data", "execute_result"):
|
|
216
|
-
outputs.append(
|
|
217
|
-
{
|
|
218
|
-
"data": {"text/plain": [content["data"].get("text/plain", "")]},
|
|
219
|
-
"execution_count": content["execution_count"],
|
|
220
|
-
"metadata": {},
|
|
221
|
-
"output_type": msg_type,
|
|
222
|
-
}
|
|
223
|
-
)
|
|
224
|
-
elif msg_type == "error":
|
|
225
|
-
outputs.append(
|
|
226
|
-
{
|
|
227
|
-
"ename": content["ename"],
|
|
228
|
-
"evalue": content["evalue"],
|
|
229
|
-
"output_type": "error",
|
|
230
|
-
"traceback": content["traceback"],
|
|
231
|
-
}
|
|
232
|
-
)
|
|
233
|
-
elif msg_type == "status" and msg["content"]["execution_state"] == "idle":
|
|
234
|
-
break
|
|
235
|
-
|
|
236
|
-
with ydoc.begin_transaction() as txn:
|
|
237
|
-
ycell.set(txn, "outputs", outputs)
|
|
@@ -1,108 +0,0 @@
|
|
|
1
|
-
import asyncio
|
|
2
|
-
import json
|
|
3
|
-
from datetime import datetime, timezone
|
|
4
|
-
from typing import Any, Dict, List, Optional
|
|
5
|
-
from uuid import uuid4
|
|
6
|
-
|
|
7
|
-
from dateutil.parser import parse as dateutil_parse
|
|
8
|
-
|
|
9
|
-
protocol_version_info = (5, 3)
|
|
10
|
-
protocol_version = "%i.%i" % protocol_version_info
|
|
11
|
-
|
|
12
|
-
|
|
13
|
-
def str_to_date(obj: Dict[str, Any]) -> Dict[str, Any]:
|
|
14
|
-
if "date" in obj:
|
|
15
|
-
obj["date"] = dateutil_parse(obj["date"])
|
|
16
|
-
return obj
|
|
17
|
-
|
|
18
|
-
|
|
19
|
-
def date_to_str(obj: Dict[str, Any]):
|
|
20
|
-
if "date" in obj and type(obj["date"]) is not str:
|
|
21
|
-
obj["date"] = obj["date"].isoformat().replace("+00:00", "Z")
|
|
22
|
-
return obj
|
|
23
|
-
|
|
24
|
-
|
|
25
|
-
def utcnow() -> datetime:
|
|
26
|
-
return datetime.utcnow().replace(tzinfo=timezone.utc)
|
|
27
|
-
|
|
28
|
-
|
|
29
|
-
def create_message_header(
|
|
30
|
-
msg_type: str, session_id: str, msg_id: str
|
|
31
|
-
) -> Dict[str, Any]:
|
|
32
|
-
if not session_id:
|
|
33
|
-
session_id = msg_id = uuid4().hex
|
|
34
|
-
else:
|
|
35
|
-
msg_id = f"{session_id}_{msg_id}"
|
|
36
|
-
header = {
|
|
37
|
-
"date": utcnow().isoformat().replace("+00:00", "Z"),
|
|
38
|
-
"msg_id": msg_id,
|
|
39
|
-
"msg_type": msg_type,
|
|
40
|
-
"session": session_id,
|
|
41
|
-
"username": "",
|
|
42
|
-
"version": protocol_version,
|
|
43
|
-
}
|
|
44
|
-
return header
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
def create_message(
|
|
48
|
-
msg_type: str,
|
|
49
|
-
content: Dict = {},
|
|
50
|
-
session_id: str = "",
|
|
51
|
-
msg_id: str = "",
|
|
52
|
-
) -> Dict[str, Any]:
|
|
53
|
-
header = create_message_header(msg_type, session_id, msg_id)
|
|
54
|
-
msg = {
|
|
55
|
-
"header": header,
|
|
56
|
-
"msg_id": header["msg_id"],
|
|
57
|
-
"msg_type": header["msg_type"],
|
|
58
|
-
"parent_header": {},
|
|
59
|
-
"content": content,
|
|
60
|
-
"metadata": {},
|
|
61
|
-
"buffers": [],
|
|
62
|
-
}
|
|
63
|
-
return msg
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
def serialize(msg: Dict[str, Any]) -> List[bytes]:
|
|
67
|
-
# TODO: buffers
|
|
68
|
-
nbufs = int(1).to_bytes(4, "big")
|
|
69
|
-
offset_0 = int(8).to_bytes(4, "big")
|
|
70
|
-
message = nbufs + offset_0 + json.dumps(msg).encode("utf8")
|
|
71
|
-
return message
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
def deserialize(
|
|
75
|
-
msg: bytes,
|
|
76
|
-
) -> Dict[str, Any]:
|
|
77
|
-
# TODO: buffers
|
|
78
|
-
# nbufs = int.from_bytes(msg[:4], byteorder="little")
|
|
79
|
-
offset_0 = int.from_bytes(msg[4:8], byteorder="little")
|
|
80
|
-
message = json.loads(msg[offset_0:])
|
|
81
|
-
return message
|
|
82
|
-
|
|
83
|
-
|
|
84
|
-
async def send_message(
|
|
85
|
-
message, channel: str, websocket, change_date_to_str: bool = False
|
|
86
|
-
):
|
|
87
|
-
_date_to_str = date_to_str if change_date_to_str else lambda x: x
|
|
88
|
-
message["header"] = _date_to_str(message["header"])
|
|
89
|
-
message["parent_header"] = _date_to_str(message["parent_header"])
|
|
90
|
-
message["metadata"] = _date_to_str(message["metadata"])
|
|
91
|
-
message["content"] = _date_to_str(message.get("content", {}))
|
|
92
|
-
message["channel"] = channel
|
|
93
|
-
# msg = serialize(message)
|
|
94
|
-
# await websocket.send_bytes(msg)
|
|
95
|
-
await websocket.send_json(message)
|
|
96
|
-
|
|
97
|
-
|
|
98
|
-
async def receive_message(
|
|
99
|
-
websocket, timeout: float, change_str_to_date: bool = False
|
|
100
|
-
) -> Optional[Dict[str, Any]]:
|
|
101
|
-
try:
|
|
102
|
-
message = await websocket.receive_json(timeout)
|
|
103
|
-
except asyncio.TimeoutError:
|
|
104
|
-
return None
|
|
105
|
-
_str_to_date = str_to_date if change_str_to_date else lambda x: x
|
|
106
|
-
message["header"] = _str_to_date(message["header"])
|
|
107
|
-
message["parent_header"] = _str_to_date(message["parent_header"])
|
|
108
|
-
# return deserialize(message, change_str_to_date)
|
|
File without changes
|
|
File without changes
|