diff --git a/backend/fluksio/api/routes/flows.py b/backend/fluksio/api/routes/flows.py index bbdd280..8a5602c 100644 --- a/backend/fluksio/api/routes/flows.py +++ b/backend/fluksio/api/routes/flows.py @@ -220,6 +220,11 @@ def _endpoints(controller: FlowController, definition: FlowDef) -> list[Endpoint def _detail(controller: FlowController, definition: FlowDef) -> FlowDetail: + """The flow and how its nodes are doing. Blocking: call it off the loop. + + A draft is built to report on it, which asks a worker to compile its + nodes, and that waits on the socket the loop serves. + """ name = definition.name running = { "enabled": controller.is_enabled(name), @@ -534,7 +539,7 @@ async def save_flow( status_code=409, detail={"message": str(exc), "current_version": exc.current}, ) - return _detail(controller, stored) + return await run_in_threadpool(_detail, controller, stored) @router.post("/{name}/publish", response_model=FlowDetail) @@ -561,7 +566,7 @@ async def publish_flow( ) _audit("published", name, user) await controller.reload_flow(name) - return _detail(controller, published) + return await run_in_threadpool(_detail, controller, published) @router.post("/{name}/discard-draft", response_model=FlowDetail) @@ -581,7 +586,7 @@ async def discard_draft(name: str, controller: FlowControllerDep) -> Any: ), ) published = await run_in_threadpool(controller.store.discard_draft, name) - return _detail(controller, published) + return await run_in_threadpool(_detail, controller, published) @router.delete("/{name}", response_model=Message) @@ -644,7 +649,7 @@ async def rename_flow( # run, so this is cleanup rather than a migration. await run_in_threadpool(controller.forget_flow, name) await controller.reload() - return _detail(controller, renamed) + return await run_in_threadpool(_detail, controller, renamed) @router.post("/{name}/duplicate", response_model=FlowDetail) @@ -662,7 +667,7 @@ async def duplicate_flow( copy = await run_in_threadpool(controller.store.duplicate_flow, name) except FlowNotFound: raise HTTPException(status_code=404, detail=f"No flow named '{name}'") - return _detail(controller, copy) + return await run_in_threadpool(_detail, controller, copy) # ----------------------------------------------------------------------------- @@ -810,7 +815,7 @@ async def share_node( ) except LibExists as exc: raise HTTPException(status_code=409, detail=str(exc)) - return _detail(controller, _read_flow(controller, name)) + return await run_in_threadpool(_detail, controller, _read_flow(controller, name)) @router.post("/{name}/nodes/{node_id}/unshare", response_model=FlowDetail) @@ -825,7 +830,7 @@ async def unshare_node( await run_in_threadpool(controller.store.unshare_node, name, node_id) except LibNotFound as exc: raise HTTPException(status_code=404, detail=str(exc)) - return _detail(controller, _read_flow(controller, name)) + return await run_in_threadpool(_detail, controller, _read_flow(controller, name)) # ----------------------------------------------------------------------------- @@ -841,7 +846,7 @@ async def start_flow( _read_flow(controller, name) await controller.set_enabled(name, True) _audit("started", name, user) - return _detail(controller, _read_flow(controller, name)) + return await run_in_threadpool(_detail, controller, _read_flow(controller, name)) @router.post("/{name}/stop", response_model=FlowDetail) @@ -850,7 +855,7 @@ async def stop_flow(name: str, controller: FlowControllerDep, user: CurrentUser) _read_flow(controller, name) await controller.set_enabled(name, False) _audit("stopped", name, user) - return _detail(controller, _read_flow(controller, name)) + return await run_in_threadpool(_detail, controller, _read_flow(controller, name)) @router.post("/{name}/pause", response_model=Message) diff --git a/backend/fluksio/api/routes/workers.py b/backend/fluksio/api/routes/workers.py index 3dca9d0..9c06631 100644 --- a/backend/fluksio/api/routes/workers.py +++ b/backend/fluksio/api/routes/workers.py @@ -337,6 +337,10 @@ async def _attached_session(websocket: WebSocket, claims: dict[str, Any]) -> Non "version": hello.get("version"), }, inventory=WorkerInventory.from_hello(hello.get("inventory")), + host=str( + hello.get("service_host") + or (websocket.client.host if websocket.client else "") + ), ) hub.attach(worker) _behind.pop(name, None) diff --git a/backend/fluksio/flow/controller.py b/backend/fluksio/flow/controller.py index ee8fdb3..6a37aa7 100644 --- a/backend/fluksio/flow/controller.py +++ b/backend/fluksio/flow/controller.py @@ -56,6 +56,7 @@ from fluksio.flow.nodes import ( Node, NtfyNode, RbeNode, + ServiceNode, SwitchNode, WebSocketNode, WrenNode, @@ -425,6 +426,13 @@ NODE_TYPES: dict[str, NodeType] = { runs_code=True, params_schema=_schema_of(ExecNode), ), + "service": NodeType( + title="Service", + description="Keep a server running on a worker, and publish its url.", + cls=ServiceNode, + runs_code=True, + params_schema=_schema_of(ServiceNode), + ), "file": NodeType( title="File", description="Read a file into the flow, or write one out of it.", diff --git a/backend/fluksio/flow/nodes/__init__.py b/backend/fluksio/flow/nodes/__init__.py index 13bda89..3cef5e0 100644 --- a/backend/fluksio/flow/nodes/__init__.py +++ b/backend/fluksio/flow/nodes/__init__.py @@ -17,6 +17,7 @@ from fluksio.flow.nodes.logic import ChangeNode, JoinNode, RbeNode, SwitchNode from fluksio.flow.nodes.mlp import MLPNode from fluksio.flow.nodes.mqtt import MqttNode from fluksio.flow.nodes.ntfy import NtfyNode +from fluksio.flow.nodes.service import ServiceNode from fluksio.flow.nodes.websocket import WebSocketNode from fluksio.flow.nodes.wren import WrenNode @@ -37,6 +38,7 @@ __all__ = [ "NtfyNode", "RESERVED_SETTINGS", "RbeNode", + "ServiceNode", "SwitchNode", "WebSocketNode", "WrenNode", diff --git a/backend/fluksio/flow/nodes/service.py b/backend/fluksio/flow/nodes/service.py new file mode 100644 index 0000000..9b2263c --- /dev/null +++ b/backend/fluksio/flow/nodes/service.py @@ -0,0 +1,246 @@ +"""Service: a server kept running on a worker while its flow is live. + +A Python node is a call that returns, in a process the worker kills afterwards, +so a model it loads is loaded again on every message. A service is the other +shape: one long call that runs a command, such as a speech server, and keeps +it up. Other nodes reach it over HTTP at the url it publishes. +""" + +from __future__ import annotations + +import asyncio +import logging +import shlex +import uuid +from collections.abc import Iterable +from typing import TYPE_CHECKING, Any + +import httpx +from pydantic import BaseModel, ConfigDict, Field + +from fluksio.core.config import settings +from fluksio.flow.messages import MessageSpec +from fluksio.flow.nodes.base import Node + +if TYPE_CHECKING: + from fastapi import FastAPI + + from fluksio.flow.remote import RemoteWorker, RemoteWorkerHub + +logger = logging.getLogger(__name__) + +#: Seconds between restarts, one step per start in a row that never got +#: healthy. A service that served and then died starts again after the first. +BACKOFF = settings.FLOW_RESTART_BACKOFF +#: Starts in a row that never answered healthy before the node reads down. +DOWN_AFTER = 3 +#: How often a starting server's health is asked, and how long one ask may take. +HEALTH_POLL_S = 1.0 +HEALTH_TIMEOUT_S = 2.0 +#: How often to look again when no worker carries the label. +WORKER_POLL_S = 2.0 +#: How long a stopping flow waits for the worker to say the server is gone, so +#: a flow started again finds the port free. Under the supervisor's grace. +STOP_WAIT_S = 4.0 + + +class ServiceNode(Node): + """Keep a server running on a worker and publish the url it listens on. + + The command runs on a worker matched by label, in one call that takes none + of the worker's ``--parallel`` slots, and starts again whenever it ends. + Every output port carries the url once the server answers its health + check, and an empty string the moment the call ends. + """ + + class Params(BaseModel): + model_config = ConfigDict(extra="allow") + + command: str = Field( + description=( + "The server to run, such as `parrot serve --host 0.0.0.0 " + "--port {port}`. {port} and {weights} are filled in." + ) + ) + port: int = Field( + default=8100, ge=1, le=65535, description="The port the server listens on." + ) + worker: str = Field( + description="The label of the worker to run it on, such as gpu." + ) + health: str = Field( + default="/health", + description="The path that answers 200 once the server is ready.", + ) + env: dict[str, str] = Field( + default_factory=dict, + description="Environment variables for the server.", + ) + weights: str = Field( + default="", + pattern=r"^(sha256:[0-9a-f]{64})?$", + description=( + "An artifact digest, fetched to the worker. Its local path " + "is {weights} in the command." + ), + ) + + __slots__ = ("cfg", "argv", "hub", "_task", "_url") + + def __init__( + self, + requires: MessageSpec | Iterable[MessageSpec] = (), + provides: MessageSpec | Iterable[MessageSpec] = (), + params: dict[str, Any] | None = None, + name: str | None = None, + ): + self.cfg = self.Params.model_validate(params or {}) + self.argv = shlex.split(self.cfg.command.replace("{port}", str(self.cfg.port))) + if not self.argv: + raise ValueError("a service needs a command") + if "{weights}" in self.cfg.command and not self.cfg.weights: + raise ValueError("the command uses {weights}, but no weights are set") + # Set from the app on start; a test sets it directly. + self.hub: RemoteWorkerHub | None = None + self._task: asyncio.Task[None] | None = None + self._url = "" + super().__init__( + f=self._never_called, + requires=requires, + provides=provides, + params=params, + name=name or "service", + ) + + @staticmethod + def _never_called(params: dict[str, Any], **kwargs: Any) -> dict[str, Any] | None: + """Its flow starts it, not a message, and it publishes on its own.""" + return None + + # -- lifecycle ------------------------------------------------------------ + + async def start(self, app: FastAPI | None = None) -> None: + if app is not None: + self.hub = getattr(app.state, "worker_hub", None) + # Spawned on every start, as the websocket node's: the supervisor + # ignores a duplicate. + if self._task is not None and not self._task.done(): + return + self._task = self._run_supervised("service", self._keep_running) + + async def stop(self, app: FastAPI | None = None) -> None: + if self._task is not None: + await self._cancel_task(self._task) + self._task = None + + async def _keep_running(self) -> None: + """Run the server, and run it again whenever it ends.""" + failures = 0 + while True: + worker = self.hub.pick(self.cfg.worker) if self.hub is not None else None + if worker is None: + self.report_health( + "degraded", f"no worker labelled '{self.cfg.worker}' is attached" + ) + await asyncio.sleep(WORKER_POLL_S) + continue + self.report_health("degraded", f"starting on '{worker.name}'") + healthy, reason = await self._serve_once(worker) + # Only starts that never got healthy lead to down: one that served + # for a while is restarted as often as it dies. + failures = 0 if healthy else failures + 1 + delay = BACKOFF[min(failures, len(BACKOFF) - 1)] + logger.warning("Service '%s' ended: %s", self.id, reason) + self.report_health( + "down" if failures >= DOWN_AFTER else "degraded", + f"{reason}; restarting in {delay:g}s", + ) + await asyncio.sleep(delay) + + async def _serve_once(self, worker: RemoteWorker) -> tuple[bool, str]: + """One call on this worker, until it ends: whether it got healthy, and why.""" + call_id = f"{self.id}:{uuid.uuid4().hex[:8]}" + host = f"[{worker.host}]" if ":" in worker.host else worker.host + url = f"http://{host}:{self.cfg.port}" + payload = { + "op": "serve", + "call_id": call_id, + "argv": self.argv, + "env": self.cfg.env, + "weights": self.cfg.weights, + } + # On a thread, as every call is, and holding no slot. + call = asyncio.ensure_future( + asyncio.to_thread(worker.request, payload, 0, slot=False) + ) + # Retrieved however it ends, so one that outlives a stop is not + # reported at exit. + call.add_done_callback(lambda done: done.cancelled() or done.exception()) + healthy = False + try: + healthy = await self._await_health(call, worker.name, url) + if healthy: + self._url = url + await asyncio.to_thread(self._publish, url) + self.report_health("ok") + await asyncio.wait({call}) + error = call.result().get("error") or {} + return healthy, str(error.get("message") or "the service stopped") + except Exception as exc: + # The worker went away or went silent, or the url was refused. + return healthy, f"{type(exc).__name__}: {exc}" + finally: + if not call.done(): + # The flow stopped, or this side failed. Ending the call kills + # the server's process group on the worker, and its answer says + # that is done, so the next start finds the port free. + await asyncio.to_thread(worker.cancel, call_id) + await asyncio.wait({call}, timeout=STOP_WAIT_S) + await asyncio.to_thread(self._withdraw) + + async def _await_health( + self, call: asyncio.Future[Any], worker: str, url: str + ) -> bool: + """Ask the server's health until it answers 200, or the call ends first.""" + async with httpx.AsyncClient(timeout=HEALTH_TIMEOUT_S) as client: + while not call.done(): + problem = await self._probe(client, url + self.cfg.health) + if problem is None: + return True + self.report_health( + "degraded", + f"starting on '{worker}': {self.cfg.health} {problem}", + ) + await asyncio.wait({call}, timeout=HEALTH_POLL_S) + return False + + @staticmethod + async def _probe(client: httpx.AsyncClient, url: str) -> str | None: + """None once the url answers 200, else what it did instead.""" + try: + response = await client.get(url) + except httpx.HTTPError as exc: + return f"is not answering ({type(exc).__name__})" + if response.status_code == 200: + return None + return f"answers {response.status_code}" + + # -- the url ---------------------------------------------------------------- + + def _publish(self, url: str) -> None: + self.emit({spec.port: url for spec in self.output_ports}) + + def _withdraw(self) -> None: + """Empty the url the moment the server is gone. + + Unless its value is gone already: a deleted flow's values are cleared + before its nodes stop, and writing the url would bring one back. + """ + if not self._url: + return + self._url = "" + pipeline = self._pipeline + if pipeline is not None and any( + pipeline.state.get(spec.name) for spec in self.output_ports + ): + self._publish("") diff --git a/backend/fluksio/flow/remote.py b/backend/fluksio/flow/remote.py index 8b31f3b..965281d 100644 --- a/backend/fluksio/flow/remote.py +++ b/backend/fluksio/flow/remote.py @@ -45,8 +45,9 @@ SEND_TIMEOUT_S = 30.0 #: whose node has no timeout of its own. SILENCE_S = 90.0 #: Protocol version this engine speaks. A worker announcing anything else is -#: refused rather than half-understood. -PROTOCOL = 2 +#: refused rather than half-understood. 3: the ``serve`` call, and calls that +#: end with their process group. +PROTOCOL = 3 class NoWorker(RemoteError): @@ -102,9 +103,13 @@ class RemoteWorker: max_parallel: int = 1, info: dict[str, Any] | None = None, inventory: WorkerInventory | None = None, + host: str = "", ) -> None: self.name = name self.labels = set(labels) + #: Where the engine reaches what this worker serves: the address its + #: socket came from, unless the worker says otherwise. + self.host = host self.info = info or {} self.inventory = inventory or WorkerInventory() # This machine's books, held here so that a worker reconnecting — which @@ -163,14 +168,21 @@ class RemoteWorker: payload: dict[str, Any], timeout: float, on_event: Callable[[dict[str, Any]], None] | None = None, + slot: bool = True, ) -> dict[str, Any]: + """Send a call and wait for its answer, on a thread rather than the loop. + + ``slot`` off is for what must not queue behind running calls, nor hold + one: a compile, and a service, which would hold one for good. + """ if self._gone: raise RemoteError(f"worker '{self.name}' is no longer attached") call_id = str(payload["call_id"]) inbox: queue.Queue[dict[str, Any] | None] = queue.Queue() # Blocking here is the backpressure, exactly as taking a slot is in the # local pool. - self._slots.acquire() + if slot: + self._slots.acquire() with self._lock: self._pending[call_id] = inbox try: @@ -226,7 +238,8 @@ class RemoteWorker: finally: with self._lock: self._pending.pop(call_id, None) - self._slots.release() + if slot: + self._slots.release() @property def compiled(self) -> set[str]: @@ -440,6 +453,7 @@ class RemoteWorkerHub: "source": source, }, timeout=timeout, + slot=False, ) except RemoteError as exc: return f"{type(exc).__name__}: {exc}" diff --git a/backend/tests/flow/test_node_types.py b/backend/tests/flow/test_node_types.py index 62e8799..c5fadc8 100644 --- a/backend/tests/flow/test_node_types.py +++ b/backend/tests/flow/test_node_types.py @@ -85,6 +85,11 @@ FIXTURES: dict[str, dict] = { "requires": [MessageSpec(name="go", dtype=DType.BOOL)], "provides": [MessageSpec(name="stdout", dtype=DType.STR)], }, + "service": { + "params": {"command": "parrot serve --port {port}", "worker": "gpu"}, + "requires": [], + "provides": [MessageSpec(name="url", dtype=DType.STR)], + }, "file": { "params": {"path": "readings.log", "mode": "append"}, "requires": [MessageSpec(name="line", dtype=DType.STR)], diff --git a/backend/tests/flow/test_remote.py b/backend/tests/flow/test_remote.py index e872e5f..83687ba 100644 --- a/backend/tests/flow/test_remote.py +++ b/backend/tests/flow/test_remote.py @@ -243,6 +243,43 @@ def test_cancelling_a_run_reaches_only_that_run(loop): assert [frame["call_id"] for frame in cancels] == ["run-a:flow.node"] +def test_a_service_and_a_compile_leave_the_only_slot_to_calls(loop): + """A service holding the slot for good would leave the worker blocked.""" + hub = RemoteWorkerHub() + socket = FakeSocket() + worker = RemoteWorker( + name="gpu1", labels=["gpu"], send=socket.send_json, loop=loop, max_parallel=1 + ) + hub.attach(worker) + call_in_thread( + lambda: worker.request({"op": "serve", "call_id": "svc"}, 0, slot=False) + ) + assert socket.arrived.wait(5) + socket.arrived.clear() + result: dict = {} + thread = call_in_thread( + lambda: result.update( + value=hub.run("gpu", "flow", "node", "src", {}, "flow.node", timeout=5) + ) + ) + assert socket.arrived.wait(5) + # Neither waits for the slot the call above holds. + compiled = call_in_thread( + lambda: result.update(problem=hub.compile("gpu", "flow", "other", "src")) + ) + deadline = time.monotonic() + 5 + while len(socket.sent) < 3 and time.monotonic() < deadline: + time.sleep(0.01) + for frame in socket.sent[1:]: + worker.deliver({"call_id": frame["call_id"], "ok": True, "result": None}) + thread.join(timeout=5) + compiled.join(timeout=5) + + assert [frame["op"] for frame in socket.sent] == ["serve", "run", "compile"] + assert result == {"value": None, "problem": None} + assert worker.in_flight == 1 + + def test_a_cancel_after_the_loop_closed_leaves_no_coroutine_behind(): closed = asyncio.new_event_loop() closed.close() diff --git a/backend/tests/flow/test_service_node.py b/backend/tests/flow/test_service_node.py new file mode 100644 index 0000000..c2de351 --- /dev/null +++ b/backend/tests/flow/test_service_node.py @@ -0,0 +1,134 @@ +"""The service node: a server on a worker, restarted when it ends, its url following.""" + +import asyncio +from collections.abc import Callable + +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="amcomp", + labels=["gpu"], + send=agent.send_json, + loop=asyncio.get_running_loop(), + host="10.0.0.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://10.0.0.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://10.0.0.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_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="amcomp", + 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()) diff --git a/backend/tests/flow/test_worker_agent.py b/backend/tests/flow/test_worker_agent.py index e64290a..5d8d5e6 100644 --- a/backend/tests/flow/test_worker_agent.py +++ b/backend/tests/flow/test_worker_agent.py @@ -1,6 +1,10 @@ """The remote worker's own guards: where its token goes, and what node code gets.""" import argparse +import asyncio +import json +import time +from pathlib import Path import pytest from fluksio_worker import agent @@ -75,3 +79,59 @@ def test_only_a_whole_digest_is_opened(digest: str) -> None: """The digest is joined into a path, so anything else could read any file.""" with pytest.raises(ValueError, match="not an artifact reference"): _Reporter("fluksio").load_artifact({"digest": digest}) + + +def _gone(pid: int) -> bool: + """No such process, or one only waiting to be reaped.""" + try: + return Path(f"/proc/{pid}/stat").read_text().split()[2] == "Z" + except FileNotFoundError: + return True + + +def test_stopping_a_service_takes_its_process_group_with_it(tmp_path: Path) -> None: + """A child of the server left running would keep its port.""" + args = argparse.Namespace( + url="wss://engine.example.com/api/v1/workers/attach", + artifact_url="", + token="t", + ) + worker = agent.Agent(args) + child = tmp_path / "child" + sent: list[dict] = [] + + class Socket: + async def send(self, text: str) -> None: + sent.append(json.loads(text)) + + async def scenario() -> None: + script = f"sleep 300 & echo $! > {child}; wait" + request = {"op": "serve", "call_id": "svc", "argv": ["sh", "-c", script]} + task = asyncio.create_task(worker._serve(Socket(), request)) + for _ in range(500): + if child.exists() and child.read_text().strip(): + break + await asyncio.sleep(0.01) + worker._cancel("svc") + await task + + asyncio.run(scenario()) + + pid = int(child.read_text()) + deadline = time.monotonic() + 5 + while not _gone(pid) and time.monotonic() < deadline: + time.sleep(0.05) + assert _gone(pid) + assert not worker.serving + assert sent == [ + { + "call_id": "svc", + "ok": False, + "error": { + "type": "ServiceStopped", + "message": "the service was stopped", + "short": "the service was stopped", + "traceback": "", + }, + } + ] diff --git a/docs/code/nodes.md b/docs/code/nodes.md index 1332173..32803e4 100644 --- a/docs/code/nodes.md +++ b/docs/code/nodes.md @@ -335,6 +335,11 @@ the local engine when none is attached. Code is compiled on the selected machine, which needs the packages your node imports. See [Remote workers](workers.md) for setup and resource requests. +Each call runs in a fresh process that ends with the call, so a model loaded +in a node is loaded again on the next message. To keep one loaded, serve it +from a [Service node](../reference/node-types.md#service) and call it over +HTTP. + ## Sharing code between flows Promote a node's source to the shared library from its panel. Other flows can diff --git a/docs/code/workers.md b/docs/code/workers.md index ad2a643..0aec879 100644 --- a/docs/code/workers.md +++ b/docs/code/workers.md @@ -55,6 +55,7 @@ provides the same agent as `fluksio worker`. | `--python` | worker's interpreter | Python used for node code | | `--parallel` | `1` | concurrent calls | | `--artifact-url` | derived from `--url` | artifact endpoint, if hosted separately | +| `--service-host` | address the engine sees | host in a service's url, when the engine reaches this machine elsewhere, such as through a tunnel | | `--cpus` | detected | advertised CPU count | | `--gpus` | detected | advertised GPU count; explicit `0` means none | | `--ram-mb` | detected | advertised memory in MB | @@ -99,6 +100,25 @@ in-flight calls. While executing, the agent sends heartbeats every ten seconds; 90 seconds without contact fails the call. Heartbeats do not reset the node's own silence timeout. +Every call runs in a process group of its own. The agent kills the group when +the engine ends the call, when the connection drops and when the agent stops, +so nothing a node started keeps running. If the agent itself is killed, its +service manager has to clean up: systemd's default `KillMode=control-group` +does. + +## Services + +A [Service node](../reference/node-types.md#service) runs a server on a worker +for as long as its flow runs, such as a speech model that should stay loaded. +It is one long call that holds none of the worker's `--parallel` slots and no +resources, so it does not block calls or hold a GPU against a training run. +The engine checks the server's health over HTTP and restarts it whenever it +stops, and its url names the address the engine sees the worker connect from. +Set `--service-host` on the worker where that address is wrong. + +A service's `weights` setting holds an artifact digest. The worker fetches it +as node code fetches an artifact and passes the local path as `{weights}`. + ## Resources and artifacts Workers report CPU, GPU and memory inventory, and scheduling checks both labels diff --git a/docs/reference/node-types.md b/docs/reference/node-types.md index 400d08e..4afa892 100644 --- a/docs/reference/node-types.md +++ b/docs/reference/node-types.md @@ -285,6 +285,28 @@ Outputs the command's stdout, stderr and exit code. In Docker, the command runs inside the backend container. Mount host files explicitly if the command needs them. +### Service + +**`service`** — keep a server running on a [remote worker](../code/workers.md#services) +for as long as its flow runs, and publish the url it listens on. Adding one +or changing its settings requires a superuser. + +| Setting | Default | Notes | +|---|---|---| +| `command` | — | the server, such as `parrot serve --host 0.0.0.0 --port {port}` | +| `port` | `8100` | where the server listens, `{port}` in the command | +| `worker` | — | label of the worker to run it on, such as `gpu` | +| `health` | `/health` | path that answers 200 once the server is ready | +| `env` | `{}` | environment variables for the server | +| `weights` | — | an artifact digest, fetched to the worker as `{weights}` | + +The flow starts it, so it takes no inputs. Every output port carries +`http://:` once the health check answers 200, and an empty +string the moment the server stops. When it stops, it starts again after a +backoff. Its health reads `degraded` while it starts, restarts or waits for a +worker, `down` after three starts in a row that never got healthy, and `ok` +while it serves. + ### File **`file`** — read a file into the flow, or write one out of it. diff --git a/frontend/src/components/Flow/FlowNode.tsx b/frontend/src/components/Flow/FlowNode.tsx index 7dd0c71..123a976 100644 --- a/frontend/src/components/Flow/FlowNode.tsx +++ b/frontend/src/components/Flow/FlowNode.tsx @@ -16,6 +16,7 @@ import { Play, Plug, Radio, + Server, Shuffle, Split, Square, @@ -58,6 +59,7 @@ const NODE_ICONS = { join: Merge, debounce: Timer, exec: Terminal, + service: Server, file: FileText, ntfy: Bell, wren: Bird, diff --git a/worker/fluksio_worker/agent.py b/worker/fluksio_worker/agent.py index 2f79fc8..210ad4d 100644 --- a/worker/fluksio_worker/agent.py +++ b/worker/fluksio_worker/agent.py @@ -13,7 +13,12 @@ It connects *out*, so the engine needs no route back and nothing has to expose Redis. What it then does is what the engine's own worker pool does: hold a few subprocesses running ``worker_main.py``, hand each call to one, and pass back everything that comes out — including the metrics a training loop reports -while it is still running. +while it is still running. A service node's call runs a command instead, a +server kept up until the engine ends the call. + +Every call runs in a process group of its own, and ends with it: when the +engine ends the call, when the socket drops, or when this agent stops. A +process left behind would keep its port and its card. Deliberately two files with one dependency. Nothing of the engine is imported here; a worker host installs Python, ``websockets``, and whatever the nodes @@ -51,9 +56,13 @@ except ImportError: # pragma: no cover - the one dependency, named plainly log = logging.getLogger("fluksio-worker") -PROTOCOL = 2 +PROTOCOL = 3 #: Sent while a call is running, so the engine can tell working from wedged. HEARTBEAT_S = 10.0 +#: How long a stopped service gets to shut down before its group is killed. +STOP_GRACE_S = 3.0 +#: How much of what a service printed last is kept, to say why it ended. +TAIL_BYTES = 4096 #: Reconnection backs off to this and no further. MAX_BACKOFF_S = 30.0 #: Environment node code is not given; the engine's lists, see `_node_env`. @@ -84,6 +93,7 @@ class Subprocess: stdout=subprocess.PIPE, close_fds=True, env=env, + start_new_session=True, ) self._buffer = bytearray() @@ -119,8 +129,8 @@ class Subprocess: self._buffer += chunk def kill(self) -> None: - with contextlib.suppress(OSError): - self.proc.send_signal(signal.SIGKILL) + """SIGKILL to the whole group, so whatever the node started goes too.""" + _signal_group(self.proc.pid, signal.SIGKILL) with contextlib.suppress(subprocess.TimeoutExpired): self.proc.wait(timeout=5) @@ -136,10 +146,18 @@ class Agent: ) self.env["FLUKSIO_ARTIFACT_TOKEN"] = args.token self.running: dict[str, Subprocess] = {} + #: Service calls, by call id. Cancelling one stops its server. + self.serving: dict[str, asyncio.Task[None]] = {} + #: Every call of this session, so a dropped socket can end them. + self._calls: set[asyncio.Task[None]] = set() self.last_done = time.monotonic() self._idled = False async def serve_forever(self) -> None: + # A stop is a Ctrl-C: the calls end, and their process groups with them. + main = asyncio.current_task() + assert main is not None + asyncio.get_running_loop().add_signal_handler(signal.SIGTERM, main.cancel) backoff = 1.0 while True: try: @@ -175,6 +193,7 @@ class Agent: "max_parallel": self.args.parallel, "venv_digest": _venv_digest(self.args.python), "inventory": _inventory(self.args), + "service_host": self.args.service_host, } ) ) @@ -207,11 +226,10 @@ class Agent: # this import" is a question about *this* machine's # packages. log.info("%s %s", op, message.get("call_id")) - task = asyncio.create_task(self._run(socket, message)) - # Without this a failure in the task is only noticed - # when it is garbage collected, which reads as a call - # that vanished. - task.add_done_callback(_report_failure) + self._spawn(self._run(socket, message)) + elif op == "serve": + log.info("serve %s", message.get("call_id")) + self._spawn(self._serve(socket, message)) elif op == "cancel": self._cancel(str(message.get("call_id") or "")) except Exception: @@ -221,9 +239,26 @@ class Agent: finally: if watchdog is not None: watchdog.cancel() + # The engine has failed these calls already, and a process + # left running would keep its port for the next one. + await self._end_calls() if self._idled: raise Idle + def _spawn(self, call: Any) -> None: + task = asyncio.create_task(call) + self._calls.add(task) + task.add_done_callback(self._calls.discard) + # Without this a failure in the task is only noticed when it is + # garbage collected, which reads as a call that vanished. + task.add_done_callback(_report_failure) + + async def _end_calls(self) -> None: + calls = list(self._calls) + for task in calls: + task.cancel() + await asyncio.gather(*calls, return_exceptions=True) + async def _idle_watch(self, socket: Any) -> None: """Close the socket once nothing has run here for ``--max-idle``. @@ -234,7 +269,7 @@ class Agent: tick = min(30.0, max(1.0, self.args.max_idle / 4)) while True: await asyncio.sleep(tick) - if self.running: + if self.running or self.serving: continue if time.monotonic() - self.last_done >= self.args.max_idle: self._idled = True @@ -259,15 +294,7 @@ class Agent: worker = Subprocess(self.args.python, {**self.env, **extra}) self.running[call_id] = worker - async def beat() -> None: - while True: - await asyncio.sleep(HEARTBEAT_S) - with contextlib.suppress(Exception): - await socket.send( - json.dumps({"call_id": call_id, "event": "heartbeat"}) - ) - - heartbeat = asyncio.create_task(beat()) + heartbeat = asyncio.create_task(_beat(socket, call_id)) try: # A subprocess that died before it could be written to is reported # by the read below, which says so and ends the call — rather than @@ -304,11 +331,128 @@ class Agent: self.last_done = time.monotonic() worker.kill() + async def _serve(self, socket: Any, request: dict[str, Any]) -> None: + """Run one server until it exits or the call is ended, then say why. + + The engine checks its health and restarts it; this only runs it. What + it prints goes to this agent's log, and the last of it into the reply. + """ + call_id = str(request.get("call_id") or "") + task = asyncio.current_task() + assert task is not None + self.serving[call_id] = task + env = { + **self.env, + **{str(k): str(v) for k, v in (request.get("env") or {}).items()}, + } + argv = [str(part) for part in request.get("argv") or []] + heartbeat = asyncio.create_task(_beat(socket, call_id)) + proc: asyncio.subprocess.Process | None = None + tail = b"" + try: + weights = str(request.get("weights") or "") + if weights: + path = await self._fetch(weights, env) + argv = [part.replace("{weights}", path) for part in argv] + proc = await asyncio.create_subprocess_exec( + *argv, + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + env=env, + start_new_session=True, + ) + assert proc.stdout is not None + while chunk := await proc.stdout.read(65536): + sys.stderr.buffer.write(chunk) + sys.stderr.flush() + tail = (tail + chunk)[-TAIL_BYTES:] + reason = f"the service exited with code {await proc.wait()}" + last = tail.decode(errors="replace").strip().splitlines() + if last: + reason += f": {last[-1][:300]}" + except asyncio.CancelledError: + # Ended by the engine or by the socket going: stopping it is the + # job, and the reply below still goes if the socket is there. + reason = "the service was stopped" + except Exception as exc: + reason = f"{type(exc).__name__}: {exc}" + finally: + heartbeat.cancel() + self.serving.pop(call_id, None) + self.last_done = time.monotonic() + if proc is not None: + await _stop_group(proc) + log.info("%s ended: %s", call_id, reason) + with contextlib.suppress(Exception): + await socket.send( + json.dumps( + { + "call_id": call_id, + "ok": False, + "error": { + "type": "ServiceStopped", + "message": reason, + "short": reason, + "traceback": "", + }, + } + ) + ) + + async def _fetch(self, digest: str, env: dict[str, str]) -> str: + """A local path for an artifact, fetched as node code fetches one.""" + proc = await asyncio.create_subprocess_exec( + self.args.python, + str(WORKER_MAIN), + "--fetch", + digest, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + out, err = await proc.communicate() + if proc.returncode: + problem = err.decode(errors="replace").strip().splitlines() + raise RuntimeError( + f"could not fetch {digest}: {problem[-1] if problem else proc.returncode}" + ) + return out.decode().strip() + def _cancel(self, call_id: str) -> None: worker = self.running.get(call_id) if worker is not None: log.info("cancelling %s", call_id) worker.kill() + service = self.serving.get(call_id) + if service is not None: + log.info("stopping %s", call_id) + service.cancel() + + +async def _beat(socket: Any, call_id: str) -> None: + """Say this call is alive, for as long as it runs.""" + while True: + await asyncio.sleep(HEARTBEAT_S) + with contextlib.suppress(Exception): + await socket.send(json.dumps({"call_id": call_id, "event": "heartbeat"})) + + +def _signal_group(pid: int, sig: int) -> None: + """Signal a call's process group, which it leads. Gone already is fine.""" + with contextlib.suppress(OSError): + os.killpg(pid, sig) + + +async def _stop_group(proc: asyncio.subprocess.Process) -> None: + """SIGTERM to the group, a moment to shut down, then SIGKILL to the rest.""" + _signal_group(proc.pid, signal.SIGTERM) + try: + await asyncio.wait_for(proc.wait(), STOP_GRACE_S) + except TimeoutError: + pass + finally: + _signal_group(proc.pid, signal.SIGKILL) def _report_failure(task: asyncio.Task[Any]) -> None: @@ -475,6 +619,14 @@ def main(argv: list[str] | None = None) -> int: ) parser.add_argument("--parallel", type=int, default=1) parser.add_argument("--artifact-url", default="") + parser.add_argument( + "--service-host", + default="", + help=( + "the address the engine reaches this machine's services on " + "(default: the address the engine sees this worker connect from)" + ), + ) parser.add_argument( "--insecure", action="store_true", @@ -527,7 +679,7 @@ def main(argv: list[str] | None = None) -> int: agent = Agent(args) try: asyncio.run(agent.serve_forever()) - except KeyboardInterrupt: + except (KeyboardInterrupt, asyncio.CancelledError): log.info("stopping") except Refused as exc: log.error("the engine refused this worker: %s", exc) diff --git a/worker/fluksio_worker/worker_main.py b/worker/fluksio_worker/worker_main.py index c09bde3..7c85480 100644 --- a/worker/fluksio_worker/worker_main.py +++ b/worker/fluksio_worker/worker_main.py @@ -620,4 +620,11 @@ def main() -> None: if __name__ == "__main__": - main() + if sys.argv[1:2] == ["--fetch"]: + # The agent's way to a service's weights: fetched as node code fetches + # an artifact, with the same cache and credentials, and the path printed. + if not _valid_digest(sys.argv[2]): + raise SystemExit(f"not an artifact digest: {sys.argv[2]}") + sys.stdout.write(_fetch(sys.argv[2]) + "\n") + else: + main()