232 lines
7.8 KiB
Python
232 lines
7.8 KiB
Python
"""The service node: a server on a worker, restarted when it ends, its url following."""
|
|
|
|
import asyncio
|
|
import time
|
|
from collections.abc import Callable
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
from fluksio.flow import remote
|
|
from fluksio.flow.messages import DType, MessageSpec
|
|
from fluksio.flow.nodes import ServiceNode, service
|
|
from fluksio.flow.pipeline import Pipeline
|
|
from fluksio.flow.remote import RemoteWorker, RemoteWorkerHub
|
|
|
|
|
|
class Agent:
|
|
"""The worker's side of the socket: records calls, and ends one on cancel."""
|
|
|
|
def __init__(self) -> None:
|
|
self.sent: list[dict] = []
|
|
self.worker: RemoteWorker | None = None
|
|
|
|
async def send_json(self, payload: dict) -> None:
|
|
self.sent.append(payload)
|
|
if payload["op"] == "cancel":
|
|
self.end(payload["call_id"], "the service was stopped")
|
|
|
|
def end(self, call_id: str, reason: str) -> None:
|
|
assert self.worker is not None
|
|
self.worker.deliver(
|
|
{"call_id": call_id, "ok": False, "error": {"message": reason}}
|
|
)
|
|
|
|
def serves(self) -> list[dict]:
|
|
return [frame for frame in self.sent if frame["op"] == "serve"]
|
|
|
|
|
|
async def until(check: Callable[[], bool]) -> None:
|
|
for _ in range(500):
|
|
if check():
|
|
return
|
|
await asyncio.sleep(0.01)
|
|
raise AssertionError("never happened")
|
|
|
|
|
|
def test_a_service_comes_back_and_its_url_follows(monkeypatch):
|
|
monkeypatch.setattr(service, "BACKOFF", (0.01,))
|
|
|
|
async def healthy(client, url):
|
|
return None
|
|
|
|
monkeypatch.setattr(ServiceNode, "_probe", staticmethod(healthy))
|
|
|
|
async def scenario() -> None:
|
|
agent = Agent()
|
|
agent.worker = RemoteWorker(
|
|
name="gpu-worker",
|
|
labels=["gpu"],
|
|
send=agent.send_json,
|
|
loop=asyncio.get_running_loop(),
|
|
host="192.0.2.5",
|
|
)
|
|
hub = RemoteWorkerHub()
|
|
node = ServiceNode(
|
|
provides=[MessageSpec(name="url", dtype=DType.STR)],
|
|
params={"command": "serve --port {port}", "worker": "gpu", "port": 8110},
|
|
)
|
|
node.assign_flow("voice", "parrot")
|
|
pipeline = Pipeline(nodes=[node])
|
|
health: list[str] = []
|
|
node._on_health = lambda _node, status, _detail: health.append(status)
|
|
node.hub = hub
|
|
|
|
def url() -> object:
|
|
return pipeline.state.get("voice.url")
|
|
|
|
await node.start()
|
|
await asyncio.sleep(0.05)
|
|
# No worker carries the label yet.
|
|
assert health[-1] == "degraded"
|
|
assert not agent.sent
|
|
|
|
hub.attach(agent.worker)
|
|
await until(lambda: bool(agent.serves()))
|
|
await until(lambda: url() == "http://192.0.2.5:8110")
|
|
first = agent.serves()[0]
|
|
assert first["argv"] == ["serve", "--port", "8110"]
|
|
assert health[-1] == "ok"
|
|
|
|
# The server dies: the url goes at once, and a second call starts.
|
|
agent.end(first["call_id"], "the service exited with code 1")
|
|
await until(lambda: url() == "")
|
|
assert "degraded" in health[health.index("ok") :]
|
|
await until(lambda: len(agent.serves()) == 2)
|
|
await until(lambda: url() == "http://192.0.2.5:8110")
|
|
|
|
# Stopping the flow ends the call on the worker, and empties the url.
|
|
await node.stop()
|
|
second = agent.serves()[1]["call_id"]
|
|
assert {"op": "cancel", "call_id": second} in agent.sent
|
|
assert url() == ""
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_a_service_cancelled_with_its_loop_ends_its_call(monkeypatch):
|
|
"""A loop that cancels every task at once must not strand the call's thread.
|
|
|
|
The call's own future is cancelled with the rest, so it reads done while the
|
|
thread still waits for the worker, which then keeps the loop's shutdown
|
|
waiting until the silence bound.
|
|
"""
|
|
monkeypatch.setattr(remote, "SILENCE_S", 3.0)
|
|
|
|
async def healthy(client, url):
|
|
return None
|
|
|
|
monkeypatch.setattr(ServiceNode, "_probe", staticmethod(healthy))
|
|
agent = Agent()
|
|
|
|
async def scenario() -> None:
|
|
agent.worker = RemoteWorker(
|
|
name="gpu-worker",
|
|
labels=["gpu"],
|
|
send=agent.send_json,
|
|
loop=asyncio.get_running_loop(),
|
|
host="192.0.2.5",
|
|
)
|
|
hub = RemoteWorkerHub()
|
|
hub.attach(agent.worker)
|
|
node = ServiceNode(
|
|
provides=[MessageSpec(name="url", dtype=DType.STR)],
|
|
params={"command": "serve", "worker": "gpu"},
|
|
)
|
|
node.assign_flow("voice", "parrot")
|
|
pipeline = Pipeline(nodes=[node])
|
|
node.hub = hub
|
|
await node.start()
|
|
await until(lambda: bool(agent.serves()))
|
|
await until(lambda: pipeline.state.get("voice.url") == "http://192.0.2.5:8100")
|
|
# No stop: asyncio.run cancels the node's task and the call together.
|
|
|
|
started = time.monotonic()
|
|
asyncio.run(scenario())
|
|
assert time.monotonic() - started < 2
|
|
assert [frame["op"] for frame in agent.sent] == ["serve", "cancel"]
|
|
|
|
|
|
def test_a_service_that_never_gets_healthy_reads_down(monkeypatch):
|
|
monkeypatch.setattr(service, "BACKOFF", (0.01,))
|
|
|
|
async def refused(client, url):
|
|
return "is not answering"
|
|
|
|
monkeypatch.setattr(ServiceNode, "_probe", staticmethod(refused))
|
|
|
|
async def scenario() -> None:
|
|
agent = Agent()
|
|
agent.worker = RemoteWorker(
|
|
name="gpu-worker",
|
|
labels=["gpu"],
|
|
send=agent.send_json,
|
|
loop=asyncio.get_running_loop(),
|
|
)
|
|
hub = RemoteWorkerHub()
|
|
hub.attach(agent.worker)
|
|
node = ServiceNode(params={"command": "serve", "worker": "gpu"})
|
|
node.assign_flow("voice", "parrot")
|
|
health: list[tuple[str, str]] = []
|
|
node._on_health = lambda _node, status, detail: health.append((status, detail))
|
|
node.hub = hub
|
|
|
|
await node.start()
|
|
for attempt in range(1, service.DOWN_AFTER + 1):
|
|
await until(lambda n=attempt: len(agent.serves()) == n)
|
|
agent.end(agent.serves()[-1]["call_id"], "Address already in use")
|
|
await until(lambda: health[-1][0] == "down")
|
|
assert "Address already in use" in health[-1][1]
|
|
await node.stop()
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_a_stop_right_after_the_start_ends_the_serve_after_it(monkeypatch):
|
|
"""The stop's cancel went out on a thread of its own, so with the serve's
|
|
thread slow to start it reached the worker first. A worker ignores a cancel
|
|
for a call it does not know, and ran that server until its socket dropped.
|
|
"""
|
|
monkeypatch.setattr(remote, "SILENCE_S", 1.0)
|
|
monkeypatch.setattr(service, "STOP_WAIT_S", 1.0)
|
|
|
|
async def refused(client, url):
|
|
return "is not answering"
|
|
|
|
monkeypatch.setattr(ServiceNode, "_probe", staticmethod(refused))
|
|
|
|
class Busy(ThreadPoolExecutor):
|
|
"""Starts what it is given late, as a pool with every thread taken."""
|
|
|
|
busy = True
|
|
late = 0
|
|
|
|
def submit(self, fn, /, *args, **kwargs):
|
|
if not self.busy:
|
|
return super().submit(fn, *args, **kwargs)
|
|
self.late += 1
|
|
return super().submit(lambda: time.sleep(0.2) or fn(*args, **kwargs))
|
|
|
|
agent = Agent()
|
|
|
|
async def scenario() -> None:
|
|
pool = Busy()
|
|
asyncio.get_running_loop().set_default_executor(pool)
|
|
agent.worker = RemoteWorker(
|
|
name="gpu-worker",
|
|
labels=["gpu"],
|
|
send=agent.send_json,
|
|
loop=asyncio.get_running_loop(),
|
|
)
|
|
hub = RemoteWorkerHub()
|
|
hub.attach(agent.worker)
|
|
node = ServiceNode(params={"command": "serve", "worker": "gpu"})
|
|
node.assign_flow("voice", "parrot")
|
|
node.hub = hub
|
|
|
|
await node.start()
|
|
await until(lambda: pool.late == 1)
|
|
pool.busy = False
|
|
await node.stop()
|
|
|
|
asyncio.run(scenario())
|
|
assert [frame["op"] for frame in agent.sent] == ["serve", "cancel"]
|