Upgrade websockets lib to v14

compute_horde/compute_horde/transport/ws.py

@@ -44,7 +44,7 @@ def __init__(
         url: str,
         *,
         max_retries: int = 5,
-        base_retry_delay: int = 1,
+        base_retry_delay: float = 1,
         retry_jitter: float = 1,
     ):
         super().__init__(name)
@@ -53,10 +53,10 @@ def __init__(
         self.retry_jitter = retry_jitter
         self.max_retries = max_retries
         self.connect_lock = asyncio.Lock()
-        self._ws: websockets.WebSocketClientProtocol | None = None
+        self._ws: websockets.ClientConnection | None = None
 
     @property
-    def ws(self) -> websockets.WebSocketClientProtocol:
+    def ws(self) -> websockets.ClientConnection:
         if self._ws is None:
             raise RuntimeError("WebSocket connection has not been established")
         return self._ws
@@ -70,11 +70,11 @@ async def start(self) -> None:
 
     async def stop(self) -> None:
         async with self.connect_lock:
-            if self._ws and self._ws.open:
+            if self._ws and self._ws.state is websockets.State.OPEN:
                 await self._ws.close()
 
     async def connect(self):
-        if self._ws and self._ws.open:
+        if self._ws and self._ws.state is websockets.State.OPEN:
             return
 
         loop = asyncio.get_running_loop()

compute_horde/pyproject.toml

@@ -8,7 +8,7 @@ requires-python = "==3.11.*"
 dependencies = [
     "pydantic<3,>=2.3",
     "bittensor<9.0.0,>=8.3.1",
-    "websockets<14.0,>=13.0",
+    "websockets>=14.0,<15.0",
     "more-itertools>=10.2.0",
     "phx-class-registry>=4.1.0",
     "requests>=2.32.2",

compute_horde/tests/test_transport.py

@@ -12,10 +12,10 @@ class WSTestServer:
     port = 8765
 
     def __init__(self):
-        self.ws_server = None
+        self.ws_server: websockets.Server | None = None
         self.received: asyncio.Queue[str] = asyncio.Queue()
 
-    async def srv(self, ws, path):
+    async def srv(self, ws: websockets.ServerConnection):
         async for message in ws:
             await self.received.put(message)
 
@@ -34,17 +34,17 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
         await self.stop()
 
     @property
-    def connection(self) -> websockets.WebSocketServerProtocol | None:
-        assert len(self.ws_server.websockets) <= 1
+    def connection(self) -> websockets.ServerConnection | None:
+        assert len(self.ws_server.connections) <= 1
 
         try:
-            return next(iter(self.ws_server.websockets))
+            return next(iter(self.ws_server.connections))
         except StopIteration:
             return None
 
     @property
     def is_connected(self) -> bool:
-        return self.connection is not None and self.connection.open
+        return self.connection is not None and self.connection.state is websockets.State.OPEN
 
 
 @pytest_asyncio.fixture

executor/pyproject.toml

@@ -20,7 +20,7 @@ dependencies = [
     "nox==2024.10.09",
     "httpx~=0.26.0",
     "channels[daphne]==4.*",
-    "websockets==13.*",
+    "websockets>=14.0,<15.0",
     "compute-horde",
     "prometheus-client~=0.17.0",
     "django-prometheus==2.3.1",