Service node: a server kept running on a worker (openpanel V2)

- `service` node type: command, port, worker label, health path, env and
  weights (an artifact digest the worker fetches as {weights}). Its flow
  starts it; it restarts with backoff, reads degraded while starting or
  without a worker and down after three starts that never got healthy, and
  publishes http://<worker>:<port> once /health answers 200, "" when it ends.
- Worker protocol 3: a `serve` call, every call in its own process group,
  killed when the engine ends the call, the socket drops or the agent gets
  SIGTERM. `--service-host` overrides the address the engine sees.
- Services and compiles take no `--parallel` slot, and a flow's detail is
  built off the event loop, so a save no longer stalls the engine.

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)

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()