feat(workers): let a panel dial the GPU machine, so more than one person can use it
Remote workers connect outbound: the node dials the control plane, spends an
enrollment token, pins a certificate. That stays the default and is unchanged.
It is also structurally 1:1 — a worker process holds one endpoint, one pinned
certificate and one worker id — so a second person wanting the same GPU box has
to get shell access to it, repoint the start script at their own address and
restart, which disconnects whoever was using it. Sharing a GPU requires root on
it and evicts the incumbent, and no amount of UI work fixes that, because the
constraint is the shape of the connection.
This adds the other arrangement: the node listens, and any panel holding a key
connects to it, concurrently, with no shell access to the machine.
* NodeService mirrors WorkerService. Transport roles invert; message roles do
not — the node still sends WorkerMessage and the panel still sends
ServerMessage, so every state machine on both sides is untouched. Register
folds into the stream as the first exchange and reuses the existing
request/response messages rather than growing parallel ones.
* Keys are per panel, not per node. Revoking one person leaves everyone else
connected; a shared key would be revoked by nobody and leave no record of
who used it. Stored hashed, compared in constant time against every key so
the reply time is not an oracle, and the plaintext exists exactly once.
* Failed authentication is throttled per source address, so one stale
bookmark cannot lock out a different panel.
* A connection log records every attach, refusal and disconnect, and any
session can be kicked. That is what replaces per-job approval, which would
make a shared GPU unusable and train people to click yes.
* Artifacts invert too: the panel pushes inputs before assigning, and fetches
results after. The node stages both under one contained directory and
trusts no id or filename off the wire.
Runs in plaintext by deliberate decision, recorded with its accepted risk in
docs/adr/inbound-node-mode.md, and scoped there to LAN and self-hosted use —
never a fleet transport, which goal_v2 B2/B5.2 still require to dial out.
Off by default, and bound to 127.0.0.1 until someone explicitly widens it.
This commit is contained in:
velixio
@@ -0,0 +1,20 @@
|
||||
"""Inbound mode: the control plane dials the node, instead of the reverse.
|
||||
|
||||
Default remote-worker mode is outbound — the node dials the control plane, which
|
||||
is what a fleet needs (goal_v2.md B2/B5.2) and what works behind NAT with no
|
||||
open ports. This package is the opposite arrangement, for two cases outbound
|
||||
cannot serve:
|
||||
|
||||
* the node is reachable but the panel is not (the panel is the laptop);
|
||||
* several people want to share one GPU box.
|
||||
|
||||
The second is the reason this exists. Outbound is structurally 1:1 — a worker
|
||||
process holds exactly one endpoint, one pinned certificate and one worker id
|
||||
(``agent.py``), so "let a colleague use the 4090" means SSHing into the box,
|
||||
repointing it and restarting, which disconnects whoever had it. Inbound is 1:N
|
||||
by construction: the node listens once and any panel holding a key connects,
|
||||
concurrently, without shell access to the machine.
|
||||
|
||||
See ``docs/adr/inbound-node-mode.md`` for the security posture, which is
|
||||
deliberately weaker than outbound's and is scoped to LAN / self-hosted use.
|
||||
"""
|
||||
@@ -0,0 +1,187 @@
|
||||
"""Artifact staging for inbound mode, where the node cannot initiate a call.
|
||||
|
||||
Outbound moves bytes with RPCs the worker starts: it pulls inputs with
|
||||
DownloadArtifact and pushes results with UploadResult. A node that was dialled
|
||||
can do neither, so both directions are driven by the panel and the node's job
|
||||
becomes staging:
|
||||
|
||||
* inputs — the panel pushes them (PushInput) *before* sending the
|
||||
assignment, so by the time the executor asks for one it is already here;
|
||||
* results — the node writes them here and names them in TaskResult; the panel
|
||||
fetches them afterwards (FetchResult).
|
||||
|
||||
Everything lands under one directory that is resolved with the repo's existing
|
||||
containment helpers. The wire supplies task ids, attempt ids and filenames, and
|
||||
none of them are trusted: this is the same asymmetry that made B13 a real
|
||||
arbitrary-write bug on the control-plane side, and it is not going to be
|
||||
reintroduced from the other end.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
from core.path_security import UnsafePath, resolve_within, safe_filename
|
||||
from worker.protocol.gen import worker_v1_pb2 as pb
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Staged bytes are deleted once fetched, but a panel that dies mid-job leaves
|
||||
# them behind. Anything older than this is swept on the next write.
|
||||
_STALE_SECONDS = 24 * 60 * 60
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Staged:
|
||||
path: str
|
||||
sha256: str
|
||||
size_bytes: int
|
||||
created_at: float
|
||||
|
||||
|
||||
class ArtifactStore:
|
||||
"""Node-side staging for one listener. Shared across panels."""
|
||||
|
||||
def __init__(self, root: str) -> None:
|
||||
self._root = os.path.abspath(root)
|
||||
self._lock = threading.Lock()
|
||||
self._out: dict[str, _Staged] = {}
|
||||
self._in: dict[str, _Staged] = {}
|
||||
os.makedirs(self._root, exist_ok=True)
|
||||
|
||||
# ── Placement ─────────────────────────────────────────────────────────
|
||||
|
||||
def _place(self, kind: str, artifact_id: str, filename: str) -> str:
|
||||
"""Build a path under the root from wire-supplied strings, safely.
|
||||
|
||||
`artifact_id` is minted here rather than taken from the wire, and the
|
||||
filename is reduced to a bare portable name before it is joined. The
|
||||
`resolve_within` call is the belt to that braces: it also rejects a
|
||||
symlink planted inside the root, which validation of the components
|
||||
alone cannot see.
|
||||
"""
|
||||
name = safe_filename(filename) if filename else ""
|
||||
if not name:
|
||||
name = "artifact.bin"
|
||||
relative = os.path.join(kind, safe_filename(artifact_id), name)
|
||||
return str(resolve_within(self._root, relative))
|
||||
|
||||
def _sweep_locked(self, now: float) -> None:
|
||||
for index in (self._out, self._in):
|
||||
for artifact_id, staged in list(index.items()):
|
||||
if now - staged.created_at <= _STALE_SECONDS:
|
||||
continue
|
||||
index.pop(artifact_id, None)
|
||||
_remove_quietly(staged.path)
|
||||
|
||||
# ── Results: node writes, panel fetches ───────────────────────────────
|
||||
|
||||
async def publish(
|
||||
self, ref: pb.TaskRef, payload: bytes, meta: dict
|
||||
) -> pb.ArtifactRef:
|
||||
"""Stage a finished result and return the ref that names it."""
|
||||
artifact_id = uuid.uuid4().hex
|
||||
filename = str(meta.get("filename") or f"{ref.attempt_id}.wav")
|
||||
digest = hashlib.sha256(payload).hexdigest()
|
||||
path = self._place("out", artifact_id, filename)
|
||||
|
||||
def write() -> None:
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
with open(path, "wb") as handle:
|
||||
handle.write(payload)
|
||||
|
||||
await asyncio.to_thread(write)
|
||||
now = time.time()
|
||||
with self._lock:
|
||||
self._sweep_locked(now)
|
||||
self._out[artifact_id] = _Staged(
|
||||
path=path, sha256=digest, size_bytes=len(payload), created_at=now
|
||||
)
|
||||
return pb.ArtifactRef(
|
||||
artifact_id=artifact_id,
|
||||
task_id=ref.task_id,
|
||||
attempt_id=ref.attempt_id,
|
||||
filename=os.path.basename(path),
|
||||
content_type=str(meta.get("content_type") or "audio/wav"),
|
||||
size_bytes=len(payload),
|
||||
sha256=digest,
|
||||
)
|
||||
|
||||
def open_result(self, artifact_id: str) -> Optional[_Staged]:
|
||||
with self._lock:
|
||||
return self._out.get(artifact_id)
|
||||
|
||||
def result_fetched(self, artifact_id: str) -> None:
|
||||
"""Drop a result once the panel has it. The panel's ack is the commit
|
||||
point, so nothing is deleted on a partial read."""
|
||||
with self._lock:
|
||||
staged = self._out.pop(artifact_id, None)
|
||||
if staged is not None:
|
||||
_remove_quietly(staged.path)
|
||||
|
||||
# ── Inputs: panel pushes, node reads ──────────────────────────────────
|
||||
|
||||
def begin_input(self, ref: pb.ArtifactRef) -> str:
|
||||
"""Reserve a staging path for an incoming push. Returns the path."""
|
||||
artifact_id = safe_filename(ref.artifact_id or uuid.uuid4().hex)
|
||||
path = self._place("in", artifact_id, ref.filename)
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
return path
|
||||
|
||||
def commit_input(self, ref: pb.ArtifactRef, path: str, digest: str, size: int) -> None:
|
||||
now = time.time()
|
||||
with self._lock:
|
||||
self._sweep_locked(now)
|
||||
self._in[ref.artifact_id] = _Staged(
|
||||
path=path, sha256=digest, size_bytes=size, created_at=now
|
||||
)
|
||||
|
||||
async def stage_in(self, ref: pb.ArtifactRef, destination: str) -> None:
|
||||
"""Hand a previously pushed input to the executor.
|
||||
|
||||
Copied rather than moved: an attempt that is retried asks for the same
|
||||
input again, and a move would make the second attempt fail with a
|
||||
missing file that no log explains.
|
||||
"""
|
||||
with self._lock:
|
||||
staged = self._in.get(ref.artifact_id)
|
||||
if staged is None:
|
||||
raise RuntimeError(
|
||||
f"the control plane did not send input {ref.artifact_id or '(unnamed)'} "
|
||||
"before assigning this task"
|
||||
)
|
||||
await asyncio.to_thread(shutil.copyfile, staged.path, destination)
|
||||
|
||||
def forget_input(self, artifact_id: str) -> None:
|
||||
with self._lock:
|
||||
staged = self._in.pop(artifact_id, None)
|
||||
if staged is not None:
|
||||
_remove_quietly(staged.path)
|
||||
|
||||
def purge(self) -> None:
|
||||
"""Drop everything. Called when the listener stops."""
|
||||
with self._lock:
|
||||
staged = list(self._out.values()) + list(self._in.values())
|
||||
self._out.clear()
|
||||
self._in.clear()
|
||||
for item in staged:
|
||||
_remove_quietly(item.path)
|
||||
|
||||
|
||||
def _remove_quietly(path: str) -> None:
|
||||
try:
|
||||
os.remove(path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
__all__ = ["ArtifactStore", "UnsafePath"]
|
||||
@@ -0,0 +1,144 @@
|
||||
"""What the node owner can see, and what they can do about it.
|
||||
|
||||
Inbound mode has no per-job approval prompt — the enable toggle is the consent
|
||||
surface, and a prompt per job would make a shared GPU unusable. That trade only
|
||||
holds if "who is using my machine right now" is answerable at a glance and
|
||||
actable in one click. Without this, a key that leaked is invisible until the
|
||||
electricity bill.
|
||||
|
||||
Events are kept in memory with a hard cap. Persisting them would put a record
|
||||
of other people's activity on disk by default, which is a bigger promise than
|
||||
this feature needs to make.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
from collections import deque
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, Optional
|
||||
|
||||
# Enough to cover a working day of joins and drops without becoming a log file
|
||||
# nobody rotates.
|
||||
_MAX_EVENTS = 200
|
||||
|
||||
|
||||
@dataclass
|
||||
class Session:
|
||||
"""One panel currently attached to this node."""
|
||||
|
||||
session_id: str
|
||||
key_id: str
|
||||
label: str
|
||||
peer: str
|
||||
connected_at: float
|
||||
tasks_run: int = 0
|
||||
# Set when the owner kicks the session. The stream loop checks it, so a
|
||||
# disconnect that arrives mid-task ends cleanly rather than by exception.
|
||||
disconnect_requested: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class Event:
|
||||
at: float
|
||||
kind: str # connected | disconnected | rejected | kicked
|
||||
label: str = ""
|
||||
peer: str = ""
|
||||
detail: str = ""
|
||||
|
||||
|
||||
class ConnectionLog:
|
||||
def __init__(self, *, now: Optional[Callable[[], float]] = None) -> None:
|
||||
self._now = now or time.time
|
||||
self._lock = threading.Lock()
|
||||
self._sessions: dict[str, Session] = {}
|
||||
self._events: deque[Event] = deque(maxlen=_MAX_EVENTS)
|
||||
|
||||
# ── Sessions ──────────────────────────────────────────────────────────
|
||||
|
||||
def opened(self, *, session_id: str, key_id: str, label: str, peer: str) -> Session:
|
||||
session = Session(
|
||||
session_id=session_id,
|
||||
key_id=key_id,
|
||||
label=label,
|
||||
peer=peer,
|
||||
connected_at=self._now(),
|
||||
)
|
||||
with self._lock:
|
||||
self._sessions[session_id] = session
|
||||
self._events.append(
|
||||
Event(at=session.connected_at, kind="connected", label=label, peer=peer)
|
||||
)
|
||||
return session
|
||||
|
||||
def closed(self, session_id: str, *, detail: str = "") -> None:
|
||||
with self._lock:
|
||||
session = self._sessions.pop(session_id, None)
|
||||
if session is None:
|
||||
return
|
||||
self._events.append(
|
||||
Event(
|
||||
at=self._now(),
|
||||
kind="kicked" if session.disconnect_requested else "disconnected",
|
||||
label=session.label,
|
||||
peer=session.peer,
|
||||
detail=detail,
|
||||
)
|
||||
)
|
||||
|
||||
def rejected(self, *, peer: str, detail: str) -> None:
|
||||
"""A refused attempt is the event that matters most and the one a
|
||||
success-only log would omit entirely."""
|
||||
with self._lock:
|
||||
self._events.append(
|
||||
Event(at=self._now(), kind="rejected", peer=peer, detail=detail)
|
||||
)
|
||||
|
||||
def task_started(self, session_id: str) -> None:
|
||||
with self._lock:
|
||||
session = self._sessions.get(session_id)
|
||||
if session is not None:
|
||||
session.tasks_run += 1
|
||||
|
||||
def kick(self, session_id: str) -> bool:
|
||||
"""Ask a session to end. Returns False if it already went away."""
|
||||
with self._lock:
|
||||
session = self._sessions.get(session_id)
|
||||
if session is None:
|
||||
return False
|
||||
session.disconnect_requested = True
|
||||
return True
|
||||
|
||||
def disconnect_requested(self, session_id: str) -> bool:
|
||||
with self._lock:
|
||||
session = self._sessions.get(session_id)
|
||||
return session is not None and session.disconnect_requested
|
||||
|
||||
# ── Reporting ─────────────────────────────────────────────────────────
|
||||
|
||||
def snapshot(self) -> dict:
|
||||
with self._lock:
|
||||
return {
|
||||
"sessions": [
|
||||
{
|
||||
"session_id": s.session_id,
|
||||
"key_id": s.key_id,
|
||||
"label": s.label,
|
||||
"peer": s.peer,
|
||||
"connected_at": s.connected_at,
|
||||
"tasks_run": s.tasks_run,
|
||||
}
|
||||
for s in sorted(self._sessions.values(), key=lambda s: s.connected_at)
|
||||
],
|
||||
"events": [
|
||||
{
|
||||
"at": e.at,
|
||||
"kind": e.kind,
|
||||
"label": e.label,
|
||||
"peer": e.peer,
|
||||
"detail": e.detail,
|
||||
}
|
||||
for e in reversed(self._events)
|
||||
],
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
"""The single copy-pasteable string that joins a panel to a node.
|
||||
|
||||
Host, port and key are three things to transcribe and three ways to fail with
|
||||
an unhelpful error, and the failures are hard to tell apart from the outside: a
|
||||
typo'd port and a wrong key both surface as "cannot connect". Collapsing them
|
||||
into one artifact with one copy button removes the whole class.
|
||||
|
||||
ovnode://ovnode_<secret>@192.168.0.110:7444
|
||||
|
||||
Deliberately URL-shaped so it survives being pasted into a chat window, and
|
||||
deliberately not an http(s) URL so no browser or link scanner treats it as
|
||||
something to fetch — a credential in a URL that something prefetches is a
|
||||
credential in somebody's access log.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
from urllib.parse import unquote, urlsplit
|
||||
|
||||
from worker.inbound.keys import KEY_PREFIX
|
||||
|
||||
SCHEME = "ovnode"
|
||||
|
||||
# Deliberately permissive about the host (IPv4, IPv6, DNS name, .local) and
|
||||
# strict about the key, because a malformed key is the recoverable mistake and
|
||||
# an unusual-looking host usually is not a mistake at all.
|
||||
_KEY_RE = re.compile(rf"^{re.escape(KEY_PREFIX)}[A-Za-z0-9_-]{{16,128}}$")
|
||||
|
||||
|
||||
class InvalidConnectionString(ValueError):
|
||||
"""Raised with a message written for the person who pasted it."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Connection:
|
||||
host: str
|
||||
port: int
|
||||
secret: str
|
||||
|
||||
@property
|
||||
def endpoint(self) -> str:
|
||||
"""host:port, bracketing IPv6 the way gRPC's resolver expects."""
|
||||
if ":" in self.host and not self.host.startswith("["):
|
||||
return f"[{self.host}]:{self.port}"
|
||||
return f"{self.host}:{self.port}"
|
||||
|
||||
def redacted(self) -> str:
|
||||
"""Safe to log. Keeps enough of the key to tell two panels apart."""
|
||||
return f"{SCHEME}://{self.secret[: len(KEY_PREFIX) + 4]}…@{self.endpoint}"
|
||||
|
||||
|
||||
def format_connection(*, host: str, port: int, secret: str) -> str:
|
||||
if ":" in host and not host.startswith("["):
|
||||
host = f"[{host}]"
|
||||
return f"{SCHEME}://{secret}@{host}:{port}"
|
||||
|
||||
|
||||
def parse_connection(text: str) -> Connection:
|
||||
"""Parse a pasted connection string, or explain what is wrong with it."""
|
||||
cleaned = (text or "").strip()
|
||||
if not cleaned:
|
||||
raise InvalidConnectionString("Paste the connection string from the GPU machine.")
|
||||
|
||||
# A bare host:port is the most likely near-miss — someone copies the
|
||||
# address out of the UI and leaves the key behind. Naming that is far more
|
||||
# use than "invalid connection string".
|
||||
if "://" not in cleaned:
|
||||
raise InvalidConnectionString(
|
||||
"That looks like an address without a key. Copy the whole "
|
||||
f"{SCHEME}://… string from the GPU machine's Settings → Remote workers."
|
||||
)
|
||||
|
||||
parts = urlsplit(cleaned)
|
||||
if parts.scheme != SCHEME:
|
||||
raise InvalidConnectionString(
|
||||
f"Expected a {SCHEME}:// connection string, but got {parts.scheme}://."
|
||||
)
|
||||
|
||||
# urlsplit puts the credential in `username` only when an `@` is present;
|
||||
# without one the key would be silently read as the hostname and the error
|
||||
# would come out as a DNS failure.
|
||||
if parts.username is None:
|
||||
raise InvalidConnectionString(
|
||||
"That connection string has no key in it. Copy the whole string, "
|
||||
"including the part before the @."
|
||||
)
|
||||
|
||||
secret = unquote(parts.username)
|
||||
if not _KEY_RE.match(secret):
|
||||
hint = (
|
||||
" Enrollment tokens (ovw_…) are for the other direction, where the "
|
||||
"GPU machine connects to you."
|
||||
if secret.startswith("ovw_")
|
||||
else ""
|
||||
)
|
||||
raise InvalidConnectionString(f"That key is not in the expected format.{hint}")
|
||||
|
||||
try:
|
||||
host, port = parts.hostname, parts.port
|
||||
except ValueError as exc:
|
||||
# urlsplit raises rather than returning None for a non-numeric or
|
||||
# out-of-range port.
|
||||
raise InvalidConnectionString("That port number is not valid.") from exc
|
||||
if not host:
|
||||
raise InvalidConnectionString("That connection string has no address in it.")
|
||||
if not port:
|
||||
raise InvalidConnectionString(
|
||||
"That connection string has no port in it. It should end in :port."
|
||||
)
|
||||
|
||||
return Connection(host=host, port=port, secret=secret)
|
||||
|
||||
|
||||
def try_parse(text: str) -> Optional[Connection]:
|
||||
try:
|
||||
return parse_connection(text)
|
||||
except InvalidConnectionString:
|
||||
return None
|
||||
@@ -0,0 +1,233 @@
|
||||
"""Per-panel API keys for inbound mode, and the throttle that protects them.
|
||||
|
||||
One key per panel, never one key for the node. A single shared key means
|
||||
revoking one person kicks everybody and forces a re-paste on every machine, so
|
||||
in practice nobody revokes and the credential outlives the reason it was
|
||||
issued. Per-key costs nothing extra at issue time and is painful to retrofit,
|
||||
because a shared key leaves no record of who used it.
|
||||
|
||||
Keys are stored hashed. The plaintext exists exactly once, in the response to
|
||||
the issuing call, and is unrecoverable afterwards — the node cannot show a key
|
||||
again later, only replace it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import secrets
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Optional
|
||||
|
||||
from worker.identity import constant_time_equals, hash_secret
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Distinguishes an inbound panel key from the `ovw_` enrollment token used by
|
||||
# outbound mode. They are never interchangeable and the prefix makes a
|
||||
# pasted-the-wrong-one mistake diagnosable instead of just "invalid".
|
||||
KEY_PREFIX = "ovnode_"
|
||||
|
||||
# 32 bytes. The same size as the enrollment-token secret, and the reason
|
||||
# `hash_secret` may be a plain SHA-256 rather than a password KDF.
|
||||
_KEY_BYTES = 32
|
||||
|
||||
# Failed-auth throttle. A key is a bearer credential with no second factor, so
|
||||
# the only thing standing between a LAN attacker and unlimited guesses is this.
|
||||
# The window is per source address: one panel typing a stale key must not lock
|
||||
# out a different panel with a good one.
|
||||
_MAX_FAILURES = 5
|
||||
_LOCKOUT_SECONDS = 60.0
|
||||
_FAILURE_WINDOW_SECONDS = 300.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class PanelKey:
|
||||
"""One panel's admission credential. The secret itself is not in here."""
|
||||
|
||||
key_id: str
|
||||
label: str
|
||||
secret_hash: str
|
||||
created_at: float
|
||||
last_seen_at: float = 0.0
|
||||
last_seen_peer: str = ""
|
||||
revoked: bool = False
|
||||
|
||||
def public(self) -> dict:
|
||||
"""The shape the UI sees. Deliberately has no field for the secret."""
|
||||
data = asdict(self)
|
||||
data.pop("secret_hash")
|
||||
return data
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Failures:
|
||||
count: int = 0
|
||||
first_at: float = 0.0
|
||||
locked_until: float = 0.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class IssuedKey:
|
||||
"""The one and only time the plaintext exists outside the caller's hands."""
|
||||
|
||||
key: PanelKey
|
||||
secret: str
|
||||
|
||||
|
||||
class KeyStore:
|
||||
"""Thread-safe, file-backed store of per-panel keys.
|
||||
|
||||
Backed by a plain JSON file rather than the settings store because the
|
||||
settings store is read by the UI process and synced into places a
|
||||
credential hash has no business being.
|
||||
"""
|
||||
|
||||
def __init__(self, path: str, *, now: Optional[callable] = None) -> None:
|
||||
self._path = path
|
||||
self._now = now or time.time
|
||||
self._lock = threading.Lock()
|
||||
self._keys: dict[str, PanelKey] = {}
|
||||
self._failures: dict[str, _Failures] = {}
|
||||
self._load()
|
||||
|
||||
# ── Persistence ───────────────────────────────────────────────────────
|
||||
|
||||
def _load(self) -> None:
|
||||
try:
|
||||
with open(self._path, encoding="utf-8") as fh:
|
||||
raw = json.load(fh)
|
||||
except (FileNotFoundError, PermissionError):
|
||||
return
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
# A corrupt file must not take the node down, but it must also not
|
||||
# silently become "no keys configured" — that reads to the user as
|
||||
# "my keys vanished" with no cause anywhere.
|
||||
logger.error(
|
||||
"The inbound key file at %s is unreadable and was ignored. "
|
||||
"Existing panels cannot connect until a key is re-issued.",
|
||||
self._path,
|
||||
)
|
||||
return
|
||||
for entry in raw.get("keys", []):
|
||||
try:
|
||||
key = PanelKey(**entry)
|
||||
except TypeError:
|
||||
continue
|
||||
self._keys[key.key_id] = key
|
||||
|
||||
def _save_locked(self) -> None:
|
||||
directory = os.path.dirname(os.path.abspath(self._path))
|
||||
os.makedirs(directory, exist_ok=True)
|
||||
payload = json.dumps(
|
||||
{"keys": [asdict(k) for k in self._keys.values()]}, indent=2
|
||||
).encode("utf-8")
|
||||
tmp = f"{self._path}.tmp"
|
||||
# 0600 from creation, never a world-readable moment — the same idiom
|
||||
# `identity.save_worker_key` uses for the Ed25519 private key.
|
||||
fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
|
||||
try:
|
||||
os.write(fd, payload)
|
||||
finally:
|
||||
os.close(fd)
|
||||
os.replace(tmp, self._path)
|
||||
try:
|
||||
os.chmod(self._path, 0o600)
|
||||
except OSError:
|
||||
# Windows and some network filesystems do not honour POSIX modes.
|
||||
pass
|
||||
|
||||
# ── Issue and revoke ──────────────────────────────────────────────────
|
||||
|
||||
def issue(self, label: str) -> IssuedKey:
|
||||
"""Mint a key for one panel. The secret is returned exactly once."""
|
||||
secret = KEY_PREFIX + secrets.token_urlsafe(_KEY_BYTES)
|
||||
now = self._now()
|
||||
key = PanelKey(
|
||||
# Derived from the secret's hash, not from a counter: it identifies
|
||||
# the key in logs without being a second thing to store, and cannot
|
||||
# be used to reconstruct the secret.
|
||||
key_id=hash_secret(secret)[:12],
|
||||
label=label.strip() or "Panel",
|
||||
secret_hash=hash_secret(secret),
|
||||
created_at=now,
|
||||
)
|
||||
with self._lock:
|
||||
self._keys[key.key_id] = key
|
||||
self._save_locked()
|
||||
return IssuedKey(key=key, secret=secret)
|
||||
|
||||
def revoke(self, key_id: str) -> bool:
|
||||
"""Revoke one panel's key. Others keep working — that is the point."""
|
||||
with self._lock:
|
||||
key = self._keys.get(key_id)
|
||||
if key is None or key.revoked:
|
||||
return False
|
||||
key.revoked = True
|
||||
self._save_locked()
|
||||
return True
|
||||
|
||||
def list_keys(self) -> list[dict]:
|
||||
with self._lock:
|
||||
return [k.public() for k in sorted(self._keys.values(), key=lambda k: k.created_at)]
|
||||
|
||||
def any_active(self) -> bool:
|
||||
with self._lock:
|
||||
return any(not k.revoked for k in self._keys.values())
|
||||
|
||||
# ── Authentication ────────────────────────────────────────────────────
|
||||
|
||||
def locked_out(self, peer: str) -> bool:
|
||||
with self._lock:
|
||||
record = self._failures.get(peer)
|
||||
return record is not None and record.locked_until > self._now()
|
||||
|
||||
def authenticate(self, secret: str, *, peer: str = "") -> Optional[PanelKey]:
|
||||
"""Return the matching live key, or None.
|
||||
|
||||
Compares against every stored key in constant time and does not stop at
|
||||
the first match. Short-circuiting would make the reply time a function
|
||||
of how many keys are configured and which one matched — a slow oracle,
|
||||
but an oracle.
|
||||
"""
|
||||
now = self._now()
|
||||
with self._lock:
|
||||
record = self._failures.get(peer)
|
||||
if record is not None and record.locked_until > now:
|
||||
return None
|
||||
|
||||
candidate = hash_secret(secret) if secret else ""
|
||||
matched: Optional[PanelKey] = None
|
||||
for key in self._keys.values():
|
||||
if key.revoked or not candidate:
|
||||
continue
|
||||
if constant_time_equals(key.secret_hash, candidate):
|
||||
matched = key
|
||||
|
||||
if matched is None:
|
||||
self._record_failure_locked(peer, now)
|
||||
return None
|
||||
|
||||
self._failures.pop(peer, None)
|
||||
matched.last_seen_at = now
|
||||
matched.last_seen_peer = peer
|
||||
self._save_locked()
|
||||
return matched
|
||||
|
||||
def _record_failure_locked(self, peer: str, now: float) -> None:
|
||||
record = self._failures.get(peer)
|
||||
if record is None or now - record.first_at > _FAILURE_WINDOW_SECONDS:
|
||||
record = _Failures(count=0, first_at=now)
|
||||
self._failures[peer] = record
|
||||
record.count += 1
|
||||
if record.count >= _MAX_FAILURES:
|
||||
record.locked_until = now + _LOCKOUT_SECONDS
|
||||
logger.warning(
|
||||
"Refusing inbound connections from %s for %.0fs after %d failed keys.",
|
||||
peer or "an unknown address",
|
||||
_LOCKOUT_SECONDS,
|
||||
record.count,
|
||||
)
|
||||
@@ -0,0 +1,364 @@
|
||||
"""The node's inbound listener: a gRPC server hosting NodeService.
|
||||
|
||||
Off by default. Enabling it is the consent surface — there is no per-job
|
||||
approval prompt, because a prompt per job makes a shared GPU unusable and
|
||||
trains people to click yes. What replaces it is visibility: every attach,
|
||||
refusal and disconnect is in the connection log, and any session can be kicked.
|
||||
|
||||
Binds to 127.0.0.1 unless the user separately and explicitly widens it. That
|
||||
default is nearly useless on its own, which is the point: reaching a node from
|
||||
another machine should be a decision someone made, not a side effect of turning
|
||||
on a feature.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
import uuid
|
||||
from typing import Callable, Optional
|
||||
|
||||
import grpc
|
||||
|
||||
from worker.inbound.artifacts import ArtifactStore
|
||||
from worker.inbound.connection_log import ConnectionLog
|
||||
from worker.inbound.keys import KeyStore
|
||||
from worker.protocol.gen import worker_v1_pb2 as pb
|
||||
from worker.protocol.gen import worker_v1_pb2_grpc as pb_grpc
|
||||
from worker.transport.client import MAX_MESSAGE_BYTES, WorkerClient
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# The panel presents its key here. Lower-case because gRPC normalises metadata
|
||||
# keys and a mixed-case constant silently never matches.
|
||||
KEY_METADATA_KEY = "x-omnivoice-node-key"
|
||||
|
||||
DEFAULT_PORT = 7444
|
||||
DEFAULT_BIND = "127.0.0.1"
|
||||
|
||||
_FETCH_CHUNK_BYTES = 1024 * 1024
|
||||
|
||||
|
||||
def _peer_of(context) -> str:
|
||||
"""A loggable source address. gRPC formats these as ipv4:1.2.3.4:5678."""
|
||||
try:
|
||||
raw = context.peer() or ""
|
||||
except Exception:
|
||||
return ""
|
||||
for prefix in ("ipv4:", "ipv6:"):
|
||||
if raw.startswith(prefix):
|
||||
return raw[len(prefix) :]
|
||||
return raw
|
||||
|
||||
|
||||
class NodeServicer(pb_grpc.NodeServiceServicer):
|
||||
"""Serves one node to any number of panels."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
keys: KeyStore,
|
||||
log: ConnectionLog,
|
||||
artifacts: ArtifactStore,
|
||||
client_factory: Callable[[ArtifactStore], WorkerClient],
|
||||
) -> None:
|
||||
self._keys = keys
|
||||
self._log = log
|
||||
self._artifacts = artifacts
|
||||
self._client_factory = client_factory
|
||||
|
||||
# ── Admission ─────────────────────────────────────────────────────────
|
||||
|
||||
def _authenticate(self, context) -> Optional[tuple[str, str]]:
|
||||
"""Return (key_id, label) or None, logging the refusal either way."""
|
||||
peer = _peer_of(context)
|
||||
if self._keys.locked_out(peer):
|
||||
self._log.rejected(peer=peer, detail="too many failed keys")
|
||||
return None
|
||||
metadata = {k.lower(): v for k, v in (context.invocation_metadata() or ())}
|
||||
secret = metadata.get(KEY_METADATA_KEY, "")
|
||||
key = self._keys.authenticate(secret, peer=peer)
|
||||
if key is None:
|
||||
self._log.rejected(
|
||||
peer=peer, detail="no key" if not secret else "key not recognised"
|
||||
)
|
||||
return None
|
||||
return key.key_id, key.label
|
||||
|
||||
# ── Attach ────────────────────────────────────────────────────────────
|
||||
|
||||
async def Attach(self, request_iterator, context):
|
||||
admitted = self._authenticate(context)
|
||||
if admitted is None:
|
||||
await context.abort(
|
||||
grpc.StatusCode.UNAUTHENTICATED,
|
||||
"This GPU machine did not recognise that key.",
|
||||
)
|
||||
return
|
||||
key_id, label = admitted
|
||||
|
||||
session_id = uuid.uuid4().hex
|
||||
peer = _peer_of(context)
|
||||
self._log.opened(session_id=session_id, key_id=key_id, label=label, peer=peer)
|
||||
client = self._client_factory(self._artifacts)
|
||||
reader: Optional[asyncio.Task] = None
|
||||
try:
|
||||
# The node speaks first even though the panel dialled: it is still
|
||||
# the side with capabilities to declare, and the panel cannot
|
||||
# schedule anything until it knows them.
|
||||
yield pb.WorkerMessage(register=client.build_register_request())
|
||||
|
||||
first = await _next_frame(request_iterator)
|
||||
if first is None or first.WhichOneof("payload") != "registered":
|
||||
await context.abort(
|
||||
grpc.StatusCode.FAILED_PRECONDITION,
|
||||
"Expected the control plane to answer with a registration.",
|
||||
)
|
||||
return
|
||||
await client.accept_registration(first.registered)
|
||||
|
||||
reader = asyncio.create_task(
|
||||
self._pump_incoming(client, request_iterator, session_id)
|
||||
)
|
||||
while True:
|
||||
if self._log.disconnect_requested(session_id):
|
||||
yield pb.WorkerMessage(
|
||||
goodbye=pb.WorkerGoodbye(
|
||||
reason="The owner of this GPU machine ended the session."
|
||||
)
|
||||
)
|
||||
return
|
||||
# Bounded so a kick lands within a second even on an idle
|
||||
# session, where nothing else would wake this loop.
|
||||
try:
|
||||
frame = await asyncio.wait_for(client.next_outbound(), timeout=1.0)
|
||||
except asyncio.TimeoutError:
|
||||
continue
|
||||
yield frame
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.warning("Inbound session from %s ended: %s", peer or "a panel", exc)
|
||||
finally:
|
||||
if reader is not None:
|
||||
reader.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError, Exception):
|
||||
await reader
|
||||
await client.stop()
|
||||
self._log.closed(session_id)
|
||||
|
||||
async def _pump_incoming(self, client: WorkerClient, request_iterator, session_id: str) -> None:
|
||||
async for message in request_iterator:
|
||||
kind = message.WhichOneof("payload")
|
||||
if kind == "assignment":
|
||||
self._log.task_started(session_id)
|
||||
if kind == "registered":
|
||||
# A second registration on a live stream is a control plane
|
||||
# bug, not a re-handshake. Ignoring it is safer than adopting
|
||||
# a new epoch mid-session and fencing the work in flight.
|
||||
logger.warning("Ignoring a repeated registration on a live session")
|
||||
continue
|
||||
await client.handle_server_message(message)
|
||||
|
||||
# ── Artifacts ─────────────────────────────────────────────────────────
|
||||
|
||||
async def FetchResult(self, request, context):
|
||||
if self._authenticate(context) is None:
|
||||
await context.abort(
|
||||
grpc.StatusCode.UNAUTHENTICATED,
|
||||
"This GPU machine did not recognise that key.",
|
||||
)
|
||||
return
|
||||
staged = self._artifacts.open_result(request.artifact_id)
|
||||
if staged is None:
|
||||
await context.abort(
|
||||
grpc.StatusCode.NOT_FOUND, "That result is no longer on this machine."
|
||||
)
|
||||
return
|
||||
|
||||
offset = int(request.size_bytes or 0)
|
||||
try:
|
||||
with open(staged.path, "rb") as handle:
|
||||
handle.seek(offset)
|
||||
while True:
|
||||
data = handle.read(_FETCH_CHUNK_BYTES)
|
||||
if not data:
|
||||
break
|
||||
chunk = pb.ResultChunk(
|
||||
ref=pb.ArtifactRef(
|
||||
artifact_id=request.artifact_id,
|
||||
task_id=request.task_id,
|
||||
attempt_id=request.attempt_id,
|
||||
filename=os.path.basename(staged.path),
|
||||
size_bytes=staged.size_bytes,
|
||||
sha256=staged.sha256,
|
||||
),
|
||||
offset=offset,
|
||||
data=data,
|
||||
)
|
||||
offset += len(data)
|
||||
chunk.last = offset >= staged.size_bytes
|
||||
yield chunk
|
||||
except OSError as exc:
|
||||
await context.abort(grpc.StatusCode.INTERNAL, f"Could not read the result: {exc}")
|
||||
return
|
||||
# Dropped only after the last chunk left this process. A panel that
|
||||
# dies mid-fetch can ask again; the stale sweep is what eventually
|
||||
# reclaims it.
|
||||
self._artifacts.result_fetched(request.artifact_id)
|
||||
|
||||
async def PushInput(self, request_iterator, context):
|
||||
if self._authenticate(context) is None:
|
||||
await context.abort(
|
||||
grpc.StatusCode.UNAUTHENTICATED,
|
||||
"This GPU machine did not recognise that key.",
|
||||
)
|
||||
return pb.ArtifactAck()
|
||||
|
||||
ref: Optional[pb.ArtifactRef] = None
|
||||
path = ""
|
||||
handle = None
|
||||
digest = hashlib.sha256()
|
||||
received = 0
|
||||
committed = False
|
||||
try:
|
||||
async for chunk in request_iterator:
|
||||
if ref is None:
|
||||
ref = chunk.ref
|
||||
path = self._artifacts.begin_input(ref)
|
||||
handle = open(path, "wb")
|
||||
if int(chunk.offset) != received:
|
||||
return pb.ArtifactAck(
|
||||
artifact_id=ref.artifact_id if ref else "",
|
||||
bytes_received=received,
|
||||
error=pb.Error(
|
||||
code="OFFSET_MISMATCH",
|
||||
message=f"expected offset {received}, got {chunk.offset}",
|
||||
),
|
||||
)
|
||||
handle.write(chunk.data)
|
||||
digest.update(chunk.data)
|
||||
received += len(chunk.data)
|
||||
if chunk.last:
|
||||
committed = True
|
||||
break
|
||||
except Exception as exc:
|
||||
return pb.ArtifactAck(
|
||||
artifact_id=ref.artifact_id if ref else "",
|
||||
bytes_received=received,
|
||||
error=pb.Error(code="INPUT_WRITE_FAILED", message=str(exc)),
|
||||
)
|
||||
finally:
|
||||
if handle is not None:
|
||||
handle.close()
|
||||
|
||||
if ref is None or not committed:
|
||||
# An iterator that simply ends is a truncated transfer, not a
|
||||
# finished one. Committing here is exactly the bug that let a
|
||||
# short upload be renamed into place and called done.
|
||||
if path:
|
||||
with contextlib.suppress(OSError):
|
||||
os.remove(path)
|
||||
return pb.ArtifactAck(
|
||||
bytes_received=received,
|
||||
error=pb.Error(
|
||||
code="INPUT_INCOMPLETE",
|
||||
message="the input ended before its final chunk",
|
||||
),
|
||||
)
|
||||
|
||||
actual = digest.hexdigest()
|
||||
if ref.sha256 and actual != ref.sha256:
|
||||
with contextlib.suppress(OSError):
|
||||
os.remove(path)
|
||||
return pb.ArtifactAck(
|
||||
artifact_id=ref.artifact_id,
|
||||
bytes_received=received,
|
||||
error=pb.Error(
|
||||
code="INPUT_CHECKSUM_MISMATCH",
|
||||
message="the input did not match the checksum the control plane declared",
|
||||
),
|
||||
)
|
||||
|
||||
self._artifacts.commit_input(ref, path, actual, received)
|
||||
return pb.ArtifactAck(
|
||||
artifact_id=ref.artifact_id, bytes_received=received, committed=True
|
||||
)
|
||||
|
||||
|
||||
async def _next_frame(request_iterator):
|
||||
try:
|
||||
return await request_iterator.__anext__()
|
||||
except StopAsyncIteration:
|
||||
return None
|
||||
|
||||
|
||||
class NodeListener:
|
||||
"""Owns the gRPC server. Started only when the user turns inbound on."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
keys: KeyStore,
|
||||
log: ConnectionLog,
|
||||
artifacts: ArtifactStore,
|
||||
client_factory: Callable[[ArtifactStore], WorkerClient],
|
||||
) -> None:
|
||||
self._servicer = NodeServicer(
|
||||
keys=keys, log=log, artifacts=artifacts, client_factory=client_factory
|
||||
)
|
||||
self._artifacts = artifacts
|
||||
self._server: Optional[grpc.aio.Server] = None
|
||||
self._bound_port = 0
|
||||
|
||||
@property
|
||||
def running(self) -> bool:
|
||||
return self._server is not None
|
||||
|
||||
@property
|
||||
def port(self) -> int:
|
||||
"""The port actually bound, which is not always the one requested."""
|
||||
return self._bound_port
|
||||
|
||||
async def start(self, *, host: str = DEFAULT_BIND, port: int = DEFAULT_PORT) -> int:
|
||||
if self._server is not None:
|
||||
return self._bound_port
|
||||
server = grpc.aio.server(
|
||||
options=[
|
||||
("grpc.max_receive_message_length", MAX_MESSAGE_BYTES),
|
||||
("grpc.max_send_message_length", MAX_MESSAGE_BYTES),
|
||||
# A panel's channel pings on an idle Attach stream exactly as a
|
||||
# worker's does outbound. Without these the server answers
|
||||
# too_many_pings and evicts the healthy panels it was waiting
|
||||
# for — the same eviction that cost this feature a day when the
|
||||
# control plane did it.
|
||||
("grpc.keepalive_permit_without_calls", 1),
|
||||
("grpc.http2.min_ping_interval_without_data_ms", 20_000),
|
||||
("grpc.http2.max_pings_without_data", 0),
|
||||
]
|
||||
)
|
||||
pb_grpc.add_NodeServiceServicer_to_server(self._servicer, server)
|
||||
# Plaintext by design; see docs/adr/inbound-node-mode.md. The API
|
||||
# key is what admits a panel, and it is not confidential against
|
||||
# anyone who can read this LAN segment.
|
||||
bind = f"[{host}]:{port}" if ":" in host and not host.startswith("[") else f"{host}:{port}"
|
||||
bound = server.add_insecure_port(bind)
|
||||
if not bound:
|
||||
raise RuntimeError(
|
||||
f"Could not listen on {bind}. Another program may already be using that port."
|
||||
)
|
||||
await server.start()
|
||||
self._server = server
|
||||
self._bound_port = bound
|
||||
logger.info("Inbound node listener accepting connections on %s", bind)
|
||||
return bound
|
||||
|
||||
async def stop(self) -> None:
|
||||
server, self._server = self._server, None
|
||||
self._bound_port = 0
|
||||
if server is not None:
|
||||
await server.stop(grace=1.0)
|
||||
self._artifacts.purge()
|
||||
File diff suppressed because one or more lines are too long
@@ -354,7 +354,7 @@ class DownloadProgress(_message.Message):
|
||||
def __init__(self, envelope: _Optional[_Union[Envelope, _Mapping]] = ..., event_json: _Optional[str] = ...) -> None: ...
|
||||
|
||||
class WorkerMessage(_message.Message):
|
||||
__slots__ = ("heartbeat", "accepted", "rejected", "model_loading", "started", "progress", "result", "failed", "cancel_ack", "capabilities", "goodbye", "pong", "download_progress")
|
||||
__slots__ = ("heartbeat", "accepted", "rejected", "model_loading", "started", "progress", "result", "failed", "cancel_ack", "capabilities", "goodbye", "pong", "download_progress", "register")
|
||||
HEARTBEAT_FIELD_NUMBER: _ClassVar[int]
|
||||
ACCEPTED_FIELD_NUMBER: _ClassVar[int]
|
||||
REJECTED_FIELD_NUMBER: _ClassVar[int]
|
||||
@@ -368,6 +368,7 @@ class WorkerMessage(_message.Message):
|
||||
GOODBYE_FIELD_NUMBER: _ClassVar[int]
|
||||
PONG_FIELD_NUMBER: _ClassVar[int]
|
||||
DOWNLOAD_PROGRESS_FIELD_NUMBER: _ClassVar[int]
|
||||
REGISTER_FIELD_NUMBER: _ClassVar[int]
|
||||
heartbeat: Heartbeat
|
||||
accepted: TaskAccepted
|
||||
rejected: TaskRejected
|
||||
@@ -381,7 +382,8 @@ class WorkerMessage(_message.Message):
|
||||
goodbye: WorkerGoodbye
|
||||
pong: Pong
|
||||
download_progress: DownloadProgress
|
||||
def __init__(self, heartbeat: _Optional[_Union[Heartbeat, _Mapping]] = ..., accepted: _Optional[_Union[TaskAccepted, _Mapping]] = ..., rejected: _Optional[_Union[TaskRejected, _Mapping]] = ..., model_loading: _Optional[_Union[TaskModelLoading, _Mapping]] = ..., started: _Optional[_Union[TaskStarted, _Mapping]] = ..., progress: _Optional[_Union[TaskProgress, _Mapping]] = ..., result: _Optional[_Union[TaskResult, _Mapping]] = ..., failed: _Optional[_Union[TaskFailed, _Mapping]] = ..., cancel_ack: _Optional[_Union[TaskCancelAck, _Mapping]] = ..., capabilities: _Optional[_Union[CapabilityUpdate, _Mapping]] = ..., goodbye: _Optional[_Union[WorkerGoodbye, _Mapping]] = ..., pong: _Optional[_Union[Pong, _Mapping]] = ..., download_progress: _Optional[_Union[DownloadProgress, _Mapping]] = ...) -> None: ...
|
||||
register: RegisterRequest
|
||||
def __init__(self, heartbeat: _Optional[_Union[Heartbeat, _Mapping]] = ..., accepted: _Optional[_Union[TaskAccepted, _Mapping]] = ..., rejected: _Optional[_Union[TaskRejected, _Mapping]] = ..., model_loading: _Optional[_Union[TaskModelLoading, _Mapping]] = ..., started: _Optional[_Union[TaskStarted, _Mapping]] = ..., progress: _Optional[_Union[TaskProgress, _Mapping]] = ..., result: _Optional[_Union[TaskResult, _Mapping]] = ..., failed: _Optional[_Union[TaskFailed, _Mapping]] = ..., cancel_ack: _Optional[_Union[TaskCancelAck, _Mapping]] = ..., capabilities: _Optional[_Union[CapabilityUpdate, _Mapping]] = ..., goodbye: _Optional[_Union[WorkerGoodbye, _Mapping]] = ..., pong: _Optional[_Union[Pong, _Mapping]] = ..., download_progress: _Optional[_Union[DownloadProgress, _Mapping]] = ..., register: _Optional[_Union[RegisterRequest, _Mapping]] = ...) -> None: ...
|
||||
|
||||
class Deadlines(_message.Message):
|
||||
__slots__ = ("accept_seconds", "model_load_seconds", "execution_seconds", "progress_lease_seconds", "result_delivery_seconds")
|
||||
@@ -501,7 +503,7 @@ class Shutdown(_message.Message):
|
||||
def __init__(self, envelope: _Optional[_Union[Envelope, _Mapping]] = ..., reason: _Optional[str] = ...) -> None: ...
|
||||
|
||||
class ServerMessage(_message.Message):
|
||||
__slots__ = ("assignment", "cancel", "result_ack", "config", "ping", "drain", "shutdown", "prewarm")
|
||||
__slots__ = ("assignment", "cancel", "result_ack", "config", "ping", "drain", "shutdown", "prewarm", "registered")
|
||||
ASSIGNMENT_FIELD_NUMBER: _ClassVar[int]
|
||||
CANCEL_FIELD_NUMBER: _ClassVar[int]
|
||||
RESULT_ACK_FIELD_NUMBER: _ClassVar[int]
|
||||
@@ -510,6 +512,7 @@ class ServerMessage(_message.Message):
|
||||
DRAIN_FIELD_NUMBER: _ClassVar[int]
|
||||
SHUTDOWN_FIELD_NUMBER: _ClassVar[int]
|
||||
PREWARM_FIELD_NUMBER: _ClassVar[int]
|
||||
REGISTERED_FIELD_NUMBER: _ClassVar[int]
|
||||
assignment: TaskAssignment
|
||||
cancel: TaskCancel
|
||||
result_ack: ResultAckMessage
|
||||
@@ -518,7 +521,8 @@ class ServerMessage(_message.Message):
|
||||
drain: Drain
|
||||
shutdown: Shutdown
|
||||
prewarm: PrewarmRequest
|
||||
def __init__(self, assignment: _Optional[_Union[TaskAssignment, _Mapping]] = ..., cancel: _Optional[_Union[TaskCancel, _Mapping]] = ..., result_ack: _Optional[_Union[ResultAckMessage, _Mapping]] = ..., config: _Optional[_Union[ConfigUpdate, _Mapping]] = ..., ping: _Optional[_Union[Ping, _Mapping]] = ..., drain: _Optional[_Union[Drain, _Mapping]] = ..., shutdown: _Optional[_Union[Shutdown, _Mapping]] = ..., prewarm: _Optional[_Union[PrewarmRequest, _Mapping]] = ...) -> None: ...
|
||||
registered: RegisterResponse
|
||||
def __init__(self, assignment: _Optional[_Union[TaskAssignment, _Mapping]] = ..., cancel: _Optional[_Union[TaskCancel, _Mapping]] = ..., result_ack: _Optional[_Union[ResultAckMessage, _Mapping]] = ..., config: _Optional[_Union[ConfigUpdate, _Mapping]] = ..., ping: _Optional[_Union[Ping, _Mapping]] = ..., drain: _Optional[_Union[Drain, _Mapping]] = ..., shutdown: _Optional[_Union[Shutdown, _Mapping]] = ..., prewarm: _Optional[_Union[PrewarmRequest, _Mapping]] = ..., registered: _Optional[_Union[RegisterResponse, _Mapping]] = ...) -> None: ...
|
||||
|
||||
class ArtifactRef(_message.Message):
|
||||
__slots__ = ("artifact_id", "task_id", "attempt_id", "filename", "content_type", "size_bytes", "sha256", "session_token")
|
||||
@@ -577,3 +581,15 @@ class ResultAck(_message.Message):
|
||||
committed: bool
|
||||
error: Error
|
||||
def __init__(self, artifact_id: _Optional[str] = ..., bytes_received: _Optional[int] = ..., committed: _Optional[bool] = ..., error: _Optional[_Union[Error, _Mapping]] = ...) -> None: ...
|
||||
|
||||
class ArtifactAck(_message.Message):
|
||||
__slots__ = ("artifact_id", "bytes_received", "committed", "error")
|
||||
ARTIFACT_ID_FIELD_NUMBER: _ClassVar[int]
|
||||
BYTES_RECEIVED_FIELD_NUMBER: _ClassVar[int]
|
||||
COMMITTED_FIELD_NUMBER: _ClassVar[int]
|
||||
ERROR_FIELD_NUMBER: _ClassVar[int]
|
||||
artifact_id: str
|
||||
bytes_received: int
|
||||
committed: bool
|
||||
error: Error
|
||||
def __init__(self, artifact_id: _Optional[str] = ..., bytes_received: _Optional[int] = ..., committed: _Optional[bool] = ..., error: _Optional[_Union[Error, _Mapping]] = ...) -> None: ...
|
||||
|
||||
@@ -234,3 +234,217 @@ class WorkerService:
|
||||
timeout,
|
||||
metadata,
|
||||
_registered_method=True)
|
||||
|
||||
|
||||
class NodeServiceStub:
|
||||
"""Hosted by the NODE, dialled by the CONTROL PLANE — the mirror image of
|
||||
WorkerService above, for the deployment where the node cannot dial out (or
|
||||
where several panels share one GPU box; see docs/adr/inbound-node-mode.md).
|
||||
|
||||
TRANSPORT roles invert here. MESSAGE roles do NOT: the node still sends
|
||||
WorkerMessage (heartbeats, capabilities, progress, results) and the control
|
||||
plane still sends ServerMessage (assignments, cancels, acks). Every state
|
||||
machine on both sides is therefore unchanged, and that is the whole point of
|
||||
mirroring the service instead of inventing a second protocol. Read the field
|
||||
names as "what this side says", never as "who called whom".
|
||||
|
||||
This service is for LAN / self-hosted use only and is never a fleet
|
||||
transport: goal_v2.md B2/B5.2 require hosted workers to dial out, and nothing
|
||||
here relaxes that.
|
||||
"""
|
||||
|
||||
def __init__(self, channel):
|
||||
"""Constructor.
|
||||
|
||||
Args:
|
||||
channel: A grpc.Channel.
|
||||
"""
|
||||
self.Attach = channel.stream_stream(
|
||||
'/omnivoice.worker.v1.NodeService/Attach',
|
||||
request_serializer=worker__v1__pb2.ServerMessage.SerializeToString,
|
||||
response_deserializer=worker__v1__pb2.WorkerMessage.FromString,
|
||||
_registered_method=True)
|
||||
self.FetchResult = channel.unary_stream(
|
||||
'/omnivoice.worker.v1.NodeService/FetchResult',
|
||||
request_serializer=worker__v1__pb2.ArtifactRef.SerializeToString,
|
||||
response_deserializer=worker__v1__pb2.ResultChunk.FromString,
|
||||
_registered_method=True)
|
||||
self.PushInput = channel.stream_unary(
|
||||
'/omnivoice.worker.v1.NodeService/PushInput',
|
||||
request_serializer=worker__v1__pb2.ArtifactChunk.SerializeToString,
|
||||
response_deserializer=worker__v1__pb2.ArtifactAck.FromString,
|
||||
_registered_method=True)
|
||||
|
||||
|
||||
class NodeServiceServicer:
|
||||
"""Hosted by the NODE, dialled by the CONTROL PLANE — the mirror image of
|
||||
WorkerService above, for the deployment where the node cannot dial out (or
|
||||
where several panels share one GPU box; see docs/adr/inbound-node-mode.md).
|
||||
|
||||
TRANSPORT roles invert here. MESSAGE roles do NOT: the node still sends
|
||||
WorkerMessage (heartbeats, capabilities, progress, results) and the control
|
||||
plane still sends ServerMessage (assignments, cancels, acks). Every state
|
||||
machine on both sides is therefore unchanged, and that is the whole point of
|
||||
mirroring the service instead of inventing a second protocol. Read the field
|
||||
names as "what this side says", never as "who called whom".
|
||||
|
||||
This service is for LAN / self-hosted use only and is never a fleet
|
||||
transport: goal_v2.md B2/B5.2 require hosted workers to dial out, and nothing
|
||||
here relaxes that.
|
||||
"""
|
||||
|
||||
def Attach(self, request_iterator, context):
|
||||
"""The control plane opens this; the node answers. Carries the same frames
|
||||
Control does, plus Register folded in as the first exchange — a node being
|
||||
dialled cannot also expose a unary Register the way a control plane does,
|
||||
and a separate round trip would leave the stream ambiguous until it
|
||||
finished.
|
||||
|
||||
The node's first frame MUST be `register` and the panel's first frame MUST
|
||||
be `registered`. Note that the node still speaks first despite the panel
|
||||
having opened the call: it is still the side with capabilities to declare.
|
||||
"""
|
||||
context.set_code(grpc.StatusCode.UNIMPLEMENTED)
|
||||
context.set_details('Method not implemented!')
|
||||
raise NotImplementedError('Method not implemented!')
|
||||
|
||||
def FetchResult(self, request, context):
|
||||
"""Artifact out, pulled instead of pushed: mirrors UploadResult. The node has
|
||||
no way to call the panel, so the panel fetches a finished result once the
|
||||
node reports it. Resumable via ArtifactRef-scoped offsets, same as
|
||||
UploadResult.
|
||||
"""
|
||||
context.set_code(grpc.StatusCode.UNIMPLEMENTED)
|
||||
context.set_details('Method not implemented!')
|
||||
raise NotImplementedError('Method not implemented!')
|
||||
|
||||
def PushInput(self, request_iterator, context):
|
||||
"""Artifact in, pushed instead of pulled: mirrors DownloadArtifact.
|
||||
"""
|
||||
context.set_code(grpc.StatusCode.UNIMPLEMENTED)
|
||||
context.set_details('Method not implemented!')
|
||||
raise NotImplementedError('Method not implemented!')
|
||||
|
||||
|
||||
def add_NodeServiceServicer_to_server(servicer, server):
|
||||
rpc_method_handlers = {
|
||||
'Attach': grpc.stream_stream_rpc_method_handler(
|
||||
servicer.Attach,
|
||||
request_deserializer=worker__v1__pb2.ServerMessage.FromString,
|
||||
response_serializer=worker__v1__pb2.WorkerMessage.SerializeToString,
|
||||
),
|
||||
'FetchResult': grpc.unary_stream_rpc_method_handler(
|
||||
servicer.FetchResult,
|
||||
request_deserializer=worker__v1__pb2.ArtifactRef.FromString,
|
||||
response_serializer=worker__v1__pb2.ResultChunk.SerializeToString,
|
||||
),
|
||||
'PushInput': grpc.stream_unary_rpc_method_handler(
|
||||
servicer.PushInput,
|
||||
request_deserializer=worker__v1__pb2.ArtifactChunk.FromString,
|
||||
response_serializer=worker__v1__pb2.ArtifactAck.SerializeToString,
|
||||
),
|
||||
}
|
||||
generic_handler = grpc.method_handlers_generic_handler(
|
||||
'omnivoice.worker.v1.NodeService', rpc_method_handlers)
|
||||
server.add_generic_rpc_handlers((generic_handler,))
|
||||
server.add_registered_method_handlers('omnivoice.worker.v1.NodeService', rpc_method_handlers)
|
||||
|
||||
|
||||
# This class is part of an EXPERIMENTAL API.
|
||||
class NodeService:
|
||||
"""Hosted by the NODE, dialled by the CONTROL PLANE — the mirror image of
|
||||
WorkerService above, for the deployment where the node cannot dial out (or
|
||||
where several panels share one GPU box; see docs/adr/inbound-node-mode.md).
|
||||
|
||||
TRANSPORT roles invert here. MESSAGE roles do NOT: the node still sends
|
||||
WorkerMessage (heartbeats, capabilities, progress, results) and the control
|
||||
plane still sends ServerMessage (assignments, cancels, acks). Every state
|
||||
machine on both sides is therefore unchanged, and that is the whole point of
|
||||
mirroring the service instead of inventing a second protocol. Read the field
|
||||
names as "what this side says", never as "who called whom".
|
||||
|
||||
This service is for LAN / self-hosted use only and is never a fleet
|
||||
transport: goal_v2.md B2/B5.2 require hosted workers to dial out, and nothing
|
||||
here relaxes that.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def Attach(request_iterator,
|
||||
target,
|
||||
options=(),
|
||||
channel_credentials=None,
|
||||
call_credentials=None,
|
||||
insecure=False,
|
||||
compression=None,
|
||||
wait_for_ready=None,
|
||||
timeout=None,
|
||||
metadata=None):
|
||||
return grpc.experimental.stream_stream(
|
||||
request_iterator,
|
||||
target,
|
||||
'/omnivoice.worker.v1.NodeService/Attach',
|
||||
worker__v1__pb2.ServerMessage.SerializeToString,
|
||||
worker__v1__pb2.WorkerMessage.FromString,
|
||||
options,
|
||||
channel_credentials,
|
||||
insecure,
|
||||
call_credentials,
|
||||
compression,
|
||||
wait_for_ready,
|
||||
timeout,
|
||||
metadata,
|
||||
_registered_method=True)
|
||||
|
||||
@staticmethod
|
||||
def FetchResult(request,
|
||||
target,
|
||||
options=(),
|
||||
channel_credentials=None,
|
||||
call_credentials=None,
|
||||
insecure=False,
|
||||
compression=None,
|
||||
wait_for_ready=None,
|
||||
timeout=None,
|
||||
metadata=None):
|
||||
return grpc.experimental.unary_stream(
|
||||
request,
|
||||
target,
|
||||
'/omnivoice.worker.v1.NodeService/FetchResult',
|
||||
worker__v1__pb2.ArtifactRef.SerializeToString,
|
||||
worker__v1__pb2.ResultChunk.FromString,
|
||||
options,
|
||||
channel_credentials,
|
||||
insecure,
|
||||
call_credentials,
|
||||
compression,
|
||||
wait_for_ready,
|
||||
timeout,
|
||||
metadata,
|
||||
_registered_method=True)
|
||||
|
||||
@staticmethod
|
||||
def PushInput(request_iterator,
|
||||
target,
|
||||
options=(),
|
||||
channel_credentials=None,
|
||||
call_credentials=None,
|
||||
insecure=False,
|
||||
compression=None,
|
||||
wait_for_ready=None,
|
||||
timeout=None,
|
||||
metadata=None):
|
||||
return grpc.experimental.stream_unary(
|
||||
request_iterator,
|
||||
target,
|
||||
'/omnivoice.worker.v1.NodeService/PushInput',
|
||||
worker__v1__pb2.ArtifactChunk.SerializeToString,
|
||||
worker__v1__pb2.ArtifactAck.FromString,
|
||||
options,
|
||||
channel_credentials,
|
||||
insecure,
|
||||
call_credentials,
|
||||
compression,
|
||||
wait_for_ready,
|
||||
timeout,
|
||||
metadata,
|
||||
_registered_method=True)
|
||||
|
||||
@@ -39,6 +39,42 @@ service WorkerService {
|
||||
rpc DownloadArtifact(ArtifactRef) returns (stream ArtifactChunk);
|
||||
}
|
||||
|
||||
// Hosted by the NODE, dialled by the CONTROL PLANE — the mirror image of
|
||||
// WorkerService above, for the deployment where the node cannot dial out (or
|
||||
// where several panels share one GPU box; see docs/adr/inbound-node-mode.md).
|
||||
//
|
||||
// TRANSPORT roles invert here. MESSAGE roles do NOT: the node still sends
|
||||
// WorkerMessage (heartbeats, capabilities, progress, results) and the control
|
||||
// plane still sends ServerMessage (assignments, cancels, acks). Every state
|
||||
// machine on both sides is therefore unchanged, and that is the whole point of
|
||||
// mirroring the service instead of inventing a second protocol. Read the field
|
||||
// names as "what this side says", never as "who called whom".
|
||||
//
|
||||
// This service is for LAN / self-hosted use only and is never a fleet
|
||||
// transport: goal_v2.md B2/B5.2 require hosted workers to dial out, and nothing
|
||||
// here relaxes that.
|
||||
service NodeService {
|
||||
// The control plane opens this; the node answers. Carries the same frames
|
||||
// Control does, plus Register folded in as the first exchange — a node being
|
||||
// dialled cannot also expose a unary Register the way a control plane does,
|
||||
// and a separate round trip would leave the stream ambiguous until it
|
||||
// finished.
|
||||
//
|
||||
// The node's first frame MUST be `register` and the panel's first frame MUST
|
||||
// be `registered`. Note that the node still speaks first despite the panel
|
||||
// having opened the call: it is still the side with capabilities to declare.
|
||||
rpc Attach(stream ServerMessage) returns (stream WorkerMessage);
|
||||
|
||||
// Artifact out, pulled instead of pushed: mirrors UploadResult. The node has
|
||||
// no way to call the panel, so the panel fetches a finished result once the
|
||||
// node reports it. Resumable via ArtifactRef-scoped offsets, same as
|
||||
// UploadResult.
|
||||
rpc FetchResult(ArtifactRef) returns (stream ResultChunk);
|
||||
|
||||
// Artifact in, pushed instead of pulled: mirrors DownloadArtifact.
|
||||
rpc PushInput(stream ArtifactChunk) returns (ArtifactAck);
|
||||
}
|
||||
|
||||
// ── Common ─────────────────────────────────────────────────────────────────
|
||||
|
||||
// Stamped on every task-scoped message so superseded work can be fenced.
|
||||
@@ -310,6 +346,16 @@ message WorkerMessage {
|
||||
WorkerGoodbye goodbye = 11;
|
||||
Pong pong = 12;
|
||||
DownloadProgress download_progress = 13;
|
||||
// Inbound mode only (NodeService.Attach): the node's opening frame, sent
|
||||
// as soon as the panel's call is authenticated. Reuses RegisterRequest
|
||||
// verbatim rather than defining a parallel message, so version
|
||||
// negotiation, capability reporting and in-flight recovery behave
|
||||
// identically in both modes — a second shape here would be a second thing
|
||||
// to keep in step forever.
|
||||
//
|
||||
// Note the direction: the node describes ITSELF, exactly as it does when
|
||||
// it dials out. Only who opened the TCP connection changed.
|
||||
RegisterRequest register = 15;
|
||||
}
|
||||
reserved 14; // future streaming frame
|
||||
}
|
||||
@@ -391,6 +437,14 @@ message ServerMessage {
|
||||
Drain drain = 6;
|
||||
Shutdown shutdown = 7;
|
||||
PrewarmRequest prewarm = 8;
|
||||
// Inbound mode only (NodeService.Attach): the panel's answer to the node's
|
||||
// `register` frame, carrying the session token and epoch. See
|
||||
// `WorkerMessage.register` for why RegisterResponse is reused as-is.
|
||||
//
|
||||
// The API key authenticating the panel travels in the call's metadata,
|
||||
// never in a frame: a credential in the stream would be copied into every
|
||||
// protocol trace and every debug log that dumps one.
|
||||
RegisterResponse registered = 9;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -428,3 +482,13 @@ message ResultAck {
|
||||
bool committed = 3;
|
||||
Error error = 4;
|
||||
}
|
||||
|
||||
// NodeService.PushInput's reply. Deliberately the same shape as ResultAck so
|
||||
// the resume logic on either side reads identically regardless of which
|
||||
// direction the bytes were travelling.
|
||||
message ArtifactAck {
|
||||
string artifact_id = 1;
|
||||
uint64 bytes_received = 2;
|
||||
bool committed = 3;
|
||||
Error error = 4;
|
||||
}
|
||||
|
||||
@@ -254,11 +254,18 @@ class WorkerClient:
|
||||
cancel: Optional[Callable[[str], Awaitable[None]]] = None,
|
||||
capability_probe: Optional[Callable[[], list[dict]]] = None,
|
||||
on_registered: Optional[Callable[[str], None]] = None,
|
||||
artifacts: Optional["ArtifactTransport"] = None,
|
||||
) -> None:
|
||||
self.config = config
|
||||
self._execute = execute
|
||||
self._cancel = cancel
|
||||
self._capability_probe = capability_probe
|
||||
# Outbound mode moves artifacts with RPCs this side initiates
|
||||
# (UploadResult / DownloadArtifact), which is only possible because
|
||||
# this side dialled. In inbound mode the node cannot call the panel at
|
||||
# all, so both directions are driven from the panel and this hook
|
||||
# swaps in the staging that makes that work. None means outbound.
|
||||
self._artifacts = artifacts
|
||||
# Lets the agent persist the server-assigned id. Without it a restarted
|
||||
# worker signs its challenge with an empty worker_id, the signature
|
||||
# never matches, and reconnecting needs a fresh enrollment token —
|
||||
@@ -323,25 +330,10 @@ class WorkerClient:
|
||||
async with self._channel() as channel:
|
||||
stub = pb_grpc.WorkerServiceStub(channel)
|
||||
response = await self._register(stub)
|
||||
if response.error.code:
|
||||
# An authentication or version refusal is not something a
|
||||
# retry loop fixes; say so rather than reconnecting forever.
|
||||
raise RuntimeError(f"{response.error.code}: {response.error.message}")
|
||||
|
||||
self._epoch = response.session_epoch
|
||||
self._session_token = response.session_token
|
||||
self.config.worker_id = response.worker_id
|
||||
# The token is spent; every later connection proves key possession.
|
||||
self.config.enrollment_token = ""
|
||||
if self._on_registered is not None:
|
||||
try:
|
||||
self._on_registered(response.worker_id)
|
||||
except Exception:
|
||||
logger.warning("Could not persist the worker id", exc_info=True)
|
||||
|
||||
authoritative = {ref.attempt_id for ref in response.authoritative_in_flight}
|
||||
await self._cancel_zombies(authoritative)
|
||||
await self._redeliver_pending()
|
||||
# An authentication or version refusal is not something a retry
|
||||
# loop fixes; `accept_registration` raises rather than reconnecting
|
||||
# forever.
|
||||
await self.accept_registration(response)
|
||||
|
||||
metadata = ((SESSION_METADATA_KEY, self._session_token),)
|
||||
stream = stub.Control(self._outbound(), metadata=metadata)
|
||||
@@ -361,7 +353,17 @@ class WorkerClient:
|
||||
# instead of the honest "no session".
|
||||
self._stub = None
|
||||
|
||||
async def _register(self, stub) -> pb.RegisterResponse:
|
||||
# ── Session seams ─────────────────────────────────────────────────────
|
||||
#
|
||||
# Outbound owns its whole connection: dial, Register, stream, repeat. A
|
||||
# node being dialled owns none of that — the gRPC servicer does — so these
|
||||
# three expose the parts that are about the PROTOCOL rather than about who
|
||||
# opened the socket. Outbound calls them through `_connect_once` exactly as
|
||||
# before; inbound calls them from the Attach handler. Neither mode gets its
|
||||
# own copy of registration, zombie reconciliation or redelivery.
|
||||
|
||||
def build_register_request(self) -> pb.RegisterRequest:
|
||||
"""This worker's self-description. Identical in both modes."""
|
||||
challenge = identity.new_challenge()
|
||||
nonce = identity.new_challenge()
|
||||
signature = self.config.keypair.sign(
|
||||
@@ -375,30 +377,58 @@ class WorkerClient:
|
||||
capabilities = (
|
||||
self._capability_probe() if self._capability_probe else self.config.capabilities
|
||||
)
|
||||
return await stub.Register(
|
||||
pb.RegisterRequest(
|
||||
envelope=pb.Envelope(sequence=self._epoch),
|
||||
protocol_version_min=PROTOCOL_VERSION,
|
||||
protocol_version_max=PROTOCOL_VERSION,
|
||||
enrollment_token=self.config.enrollment_token,
|
||||
worker_id=self.config.worker_id,
|
||||
public_key=self.config.keypair.public_bytes(),
|
||||
challenge=challenge,
|
||||
challenge_signature=signature,
|
||||
nonce=nonce,
|
||||
key_id=self.config.keypair.key_id,
|
||||
host=codec.host_to_pb(self.config.host or describe_host()),
|
||||
capabilities=[codec.capability_to_pb(c) for c in capabilities],
|
||||
max_concurrent_tasks=self.config.max_concurrent_tasks,
|
||||
in_flight=[
|
||||
codec.task_ref(t.split("/")[0], t.split("/")[1], self._epoch)
|
||||
for t in self._running
|
||||
],
|
||||
completed_unacked=[p.ref for p in self._pending.values()],
|
||||
features=sorted(REQUIRED_FEATURES),
|
||||
)
|
||||
return pb.RegisterRequest(
|
||||
envelope=pb.Envelope(sequence=self._epoch),
|
||||
protocol_version_min=PROTOCOL_VERSION,
|
||||
protocol_version_max=PROTOCOL_VERSION,
|
||||
enrollment_token=self.config.enrollment_token,
|
||||
worker_id=self.config.worker_id,
|
||||
public_key=self.config.keypair.public_bytes(),
|
||||
challenge=challenge,
|
||||
challenge_signature=signature,
|
||||
nonce=nonce,
|
||||
key_id=self.config.keypair.key_id,
|
||||
host=codec.host_to_pb(self.config.host or describe_host()),
|
||||
capabilities=[codec.capability_to_pb(c) for c in capabilities],
|
||||
max_concurrent_tasks=self.config.max_concurrent_tasks,
|
||||
in_flight=[
|
||||
codec.task_ref(t.split("/")[0], t.split("/")[1], self._epoch)
|
||||
for t in self._running
|
||||
],
|
||||
completed_unacked=[p.ref for p in self._pending.values()],
|
||||
features=sorted(REQUIRED_FEATURES),
|
||||
)
|
||||
|
||||
async def accept_registration(self, response: pb.RegisterResponse) -> None:
|
||||
"""Adopt the control plane's answer and recover in-flight state."""
|
||||
if response.error.code:
|
||||
raise RuntimeError(f"{response.error.code}: {response.error.message}")
|
||||
|
||||
self._epoch = response.session_epoch
|
||||
self._session_token = response.session_token
|
||||
self.config.worker_id = response.worker_id
|
||||
# The token is spent; every later connection proves key possession.
|
||||
self.config.enrollment_token = ""
|
||||
if self._on_registered is not None:
|
||||
try:
|
||||
self._on_registered(response.worker_id)
|
||||
except Exception:
|
||||
logger.warning("Could not persist the worker id", exc_info=True)
|
||||
|
||||
authoritative = {ref.attempt_id for ref in response.authoritative_in_flight}
|
||||
await self._cancel_zombies(authoritative)
|
||||
await self._redeliver_pending()
|
||||
|
||||
async def next_outbound(self) -> pb.WorkerMessage:
|
||||
"""The next frame this worker wants to send."""
|
||||
return await self._outbox.get()
|
||||
|
||||
async def handle_server_message(self, message: pb.ServerMessage) -> None:
|
||||
await self._on_server_message(message)
|
||||
|
||||
async def _register(self, stub) -> pb.RegisterResponse:
|
||||
return await stub.Register(self.build_register_request())
|
||||
|
||||
# ── Outbound ──────────────────────────────────────────────────────────
|
||||
|
||||
async def _outbound(self):
|
||||
@@ -749,6 +779,11 @@ class WorkerClient:
|
||||
refuse a transfer that arrives short or corrupted instead of renaming a
|
||||
truncated file into place and calling the task done.
|
||||
"""
|
||||
if self._artifacts is not None:
|
||||
# Inbound: nothing is pushed. The result is staged here and the
|
||||
# panel fetches it after the TaskResult frame names it.
|
||||
return await self._artifacts.publish(ref, payload, meta)
|
||||
|
||||
stub = self._stub
|
||||
if stub is None:
|
||||
raise RuntimeError("no session is established")
|
||||
@@ -868,6 +903,11 @@ class WorkerClient:
|
||||
|
||||
async def _fetch_input(self, ref: pb.ArtifactRef, destination: str) -> None:
|
||||
"""Download one declared input with authenticated, ordered chunks."""
|
||||
if self._artifacts is not None:
|
||||
# Inbound: the panel pushed this before it sent the assignment, so
|
||||
# there is nothing to pull — only a staged file to hand over.
|
||||
return await self._artifacts.stage_in(ref, destination)
|
||||
|
||||
if self._stub is None:
|
||||
raise RuntimeError("no session is established")
|
||||
request = pb.ArtifactRef()
|
||||
|
||||
@@ -0,0 +1,121 @@
|
||||
# Inbound node mode: the panel dials the GPU machine
|
||||
|
||||
**Status:** accepted
|
||||
**Date:** 2026-08-11
|
||||
|
||||
## Context
|
||||
|
||||
Remote workers connect *outbound*: the node dials the control plane, presents a
|
||||
single-use enrollment token, pins the control plane's certificate on first use,
|
||||
and proves possession of an Ed25519 key on every reconnect. That is the right
|
||||
default. It works behind NAT with no open ports, it is what a hosted fleet
|
||||
needs (`remote/goal_v2.md` B2), and its security posture is strong.
|
||||
|
||||
It is also structurally **1:1**. A worker process holds exactly one endpoint,
|
||||
one pinned certificate and one worker id (`backend/worker/agent.py`). So when a
|
||||
second person wants to use the same GPU box, the only route is:
|
||||
|
||||
1. get shell access to the machine — someone else's machine;
|
||||
2. mint an enrollment token on *their* panel;
|
||||
3. edit the start script to point at their address;
|
||||
4. restart, which **disconnects whoever was using it**.
|
||||
|
||||
Sharing a GPU therefore requires root on it and evicts the incumbent. For a
|
||||
household or a small team with one 4090, that is the difference between a
|
||||
feature and a thing nobody uses. No amount of polish on the outbound flow fixes
|
||||
it, because the constraint is the shape of the connection, not the UI.
|
||||
|
||||
## Decision
|
||||
|
||||
Add an **inbound** mode, alongside outbound, in which the node listens and any
|
||||
panel holding an API key connects to it. Outbound remains the default and is
|
||||
unchanged.
|
||||
|
||||
Specifics:
|
||||
|
||||
- **New gRPC service `NodeService`, hosted by the node** — `Attach`,
|
||||
`FetchResult`, `PushInput` — mirroring `WorkerService`. **Transport roles
|
||||
invert; message roles do not.** The node still sends `WorkerMessage`
|
||||
(heartbeats, capabilities, progress, results) and the panel still sends
|
||||
`ServerMessage` (assignments, cancels, acks), so every state machine on both
|
||||
sides is untouched. `Register` folds into the stream as the first exchange,
|
||||
reusing `RegisterRequest`/`RegisterResponse` verbatim rather than defining
|
||||
parallel messages.
|
||||
- **Per-panel API keys, not one node key.** Stored as SHA-256 hashes with a
|
||||
constant-time compare; the plaintext exists once, in the issuing response.
|
||||
- **Failed-auth throttle**, per source address, so one stale bookmark cannot
|
||||
lock out a different panel.
|
||||
- **Concurrent panels are allowed.** Two panels may each believe they hold the
|
||||
node's free slot and both dispatch; the worker's `WORKER_AT_CAPACITY` reject
|
||||
is authoritative (`goal_v2.md` B5.5) and the loser retries. Contention
|
||||
degrades to a queue, not to corruption. Serialising panels instead would
|
||||
recreate the eviction problem in a nicer wrapper.
|
||||
- **Ed25519 identity is retained.** The API key admits a *panel to the node*;
|
||||
the node still proves itself to the panel with its keypair. The key is
|
||||
admission, never identity.
|
||||
- **A visible connection log with a kick action** replaces per-job approval.
|
||||
- **Default bind `127.0.0.1`.** Listening on other interfaces is a separate,
|
||||
explicit setting, and the bound address is shown in the UI.
|
||||
- **Off by default.** Enabling it is the consent surface.
|
||||
|
||||
## The part we are choosing to accept: no TLS
|
||||
|
||||
Inbound runs in plaintext. The owner decided this deliberately after the risk
|
||||
was raised. Recording it here so it is a decision with a rationale rather than
|
||||
an omission someone discovers later.
|
||||
|
||||
What that means, stated plainly:
|
||||
|
||||
- The API key crosses the network in gRPC metadata **in the clear**. Anyone who
|
||||
can read the LAN segment — or sits on the path — can lift it and use the GPU.
|
||||
- There is no server authentication, so a machine that can answer at the node's
|
||||
address can impersonate it and receive whatever the panel sends, **including
|
||||
reference audio for voice cloning**.
|
||||
- Rendered audio comes back in the clear.
|
||||
|
||||
Why it was accepted: the target is a trusted LAN or a self-hosted setup where
|
||||
the alternative is people not using the feature at all. TLS here means either
|
||||
a self-signed certificate the user must pin by hand — the step that already
|
||||
makes outbound enrollment the hardest part of the flow — or a CA nobody has.
|
||||
The judgement is that the honest, documented plaintext channel is better than a
|
||||
pinning ceremony that gets abandoned halfway.
|
||||
|
||||
Consequences accepted with it:
|
||||
|
||||
- **This mode is for LAN and self-hosted use only, and is never a fleet
|
||||
transport.** `goal_v2.md` B2 and B5.2 require hosted workers to dial out with
|
||||
no inbound ports; nothing here relaxes that, and the hosted platform must not
|
||||
inherit this path. See "Amendment" below.
|
||||
- Binding to `0.0.0.0` on an untrusted network exposes a plaintext bearer
|
||||
credential to that network. The UI says so at the point the bind is widened.
|
||||
- If the threat model ever changes, TLS is additive: the listener already
|
||||
terminates its own connections, so `add_secure_port` and a pinned certificate
|
||||
slot in without touching the protocol.
|
||||
|
||||
## Amendment to goal_v2.md B5.2
|
||||
|
||||
B5.2 reads "Workers connect outbound; the control plane never dials in", and
|
||||
B2 promises "no static IP, no inbound ports". Those remain true **for the
|
||||
hosted fleet**, which is what they were written about. They are now scoped
|
||||
statements rather than global ones: the OSS desktop product also supports an
|
||||
inbound mode for locally-owned hardware, with the weaker posture documented
|
||||
above. The conformance fixtures for the Go control plane must cover
|
||||
`WorkerService` only.
|
||||
|
||||
## Alternatives rejected
|
||||
|
||||
- **Multi-endpoint dial-out** (the node dials N control planes). Preserves the
|
||||
outbound principle and solves nothing: the second user still needs SSH access
|
||||
to add their endpoint and still needs a token from their own panel. The
|
||||
complexity being removed is "you need shell access to someone else's GPU
|
||||
box", and this keeps all of it.
|
||||
- **One shared node key.** Cheaper, and revoking it kicks everyone and forces a
|
||||
re-paste on every machine — so in practice nobody revokes, and the credential
|
||||
outlives the reason it was issued. A shared key also leaves no record of who
|
||||
used it, which makes the connection log much less useful.
|
||||
- **Per-job approval prompts.** Makes a shared GPU unusable and trains people
|
||||
to click yes. Visibility plus a kick button is the better trade.
|
||||
- **Reusing `WorkerService` with the panel as gRPC client.** Does not work: the
|
||||
gRPC client sends the request-stream type, so the panel would be sending
|
||||
`WorkerMessage`. Mirroring the service is what keeps the message roles — and
|
||||
therefore every state machine — unchanged.
|
||||
@@ -0,0 +1,186 @@
|
||||
"""Admission for inbound mode: per-panel keys, hashing, throttling.
|
||||
|
||||
Inbound trades away the TLS pinning and single-use enrollment token that
|
||||
outbound relies on (docs/adr/0002-inbound-node-mode.md), so the API key is the
|
||||
whole of admission. These tests exist because everything that protects it —
|
||||
hashing at rest, per-key revocation, the failed-auth throttle — is invisible in
|
||||
normal use and would fail silently if it regressed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from worker.inbound import keys as keys_module
|
||||
from worker.inbound.connection_string import (
|
||||
InvalidConnectionString,
|
||||
format_connection,
|
||||
parse_connection,
|
||||
)
|
||||
from worker.inbound.keys import KEY_PREFIX, KeyStore
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def store(tmp_path):
|
||||
return KeyStore(str(tmp_path / "inbound-keys.json"))
|
||||
|
||||
|
||||
def test_the_plaintext_key_is_never_written_to_disk(store, tmp_path):
|
||||
"""The node can replace a key but must never be able to show it again."""
|
||||
issued = store.issue("Alice laptop")
|
||||
|
||||
on_disk = (tmp_path / "inbound-keys.json").read_text(encoding="utf-8")
|
||||
assert issued.secret not in on_disk
|
||||
assert issued.key.secret_hash in on_disk
|
||||
|
||||
# And nothing in the API hands it back either — a "reveal key" button is
|
||||
# the feature this shape exists to make impossible to build by accident.
|
||||
assert all("secret" not in row for row in store.list_keys())
|
||||
|
||||
|
||||
def test_revoking_one_panel_leaves_the_others_working(store):
|
||||
"""The whole reason keys are per-panel rather than one shared node key."""
|
||||
alice = store.issue("Alice")
|
||||
bob = store.issue("Bob")
|
||||
|
||||
assert store.revoke(alice.key.key_id) is True
|
||||
|
||||
assert store.authenticate(alice.secret, peer="10.0.0.1") is None
|
||||
assert store.authenticate(bob.secret, peer="10.0.0.2") is not None
|
||||
|
||||
|
||||
def test_a_wrong_key_is_throttled_before_it_can_be_guessed(store, monkeypatch):
|
||||
"""A bearer credential with no second factor has only this between it and
|
||||
unlimited LAN guesses."""
|
||||
store.issue("Alice")
|
||||
|
||||
for _ in range(keys_module._MAX_FAILURES):
|
||||
assert store.authenticate("ovnode_wrong", peer="10.0.0.9") is None
|
||||
|
||||
assert store.locked_out("10.0.0.9") is True
|
||||
|
||||
|
||||
def test_one_panel_typing_a_stale_key_cannot_lock_out_another(store):
|
||||
"""The throttle is per source address on purpose: a shared counter turns
|
||||
one person's stale bookmark into an outage for everybody else."""
|
||||
good = store.issue("Bob")
|
||||
|
||||
for _ in range(keys_module._MAX_FAILURES + 2):
|
||||
store.authenticate("ovnode_wrong", peer="10.0.0.9")
|
||||
|
||||
assert store.locked_out("10.0.0.9") is True
|
||||
assert store.locked_out("10.0.0.10") is False
|
||||
assert store.authenticate(good.secret, peer="10.0.0.10") is not None
|
||||
|
||||
|
||||
def test_a_locked_out_peer_is_refused_even_with_the_right_key(store):
|
||||
"""Otherwise the throttle is decorative: an attacker who eventually
|
||||
guesses correctly is admitted on the guess that succeeds."""
|
||||
good = store.issue("Bob")
|
||||
for _ in range(keys_module._MAX_FAILURES):
|
||||
store.authenticate("ovnode_wrong", peer="10.0.0.9")
|
||||
|
||||
assert store.authenticate(good.secret, peer="10.0.0.9") is None
|
||||
|
||||
|
||||
def test_an_empty_key_never_authenticates(store):
|
||||
"""A missing metadata header arrives as "" and must not match a key whose
|
||||
hash happens to be falsy-adjacent."""
|
||||
store.issue("Alice")
|
||||
assert store.authenticate("", peer="10.0.0.1") is None
|
||||
|
||||
|
||||
def test_keys_survive_a_restart(store, tmp_path):
|
||||
issued = store.issue("Alice")
|
||||
|
||||
reopened = KeyStore(str(tmp_path / "inbound-keys.json"))
|
||||
|
||||
assert reopened.authenticate(issued.secret, peer="10.0.0.1") is not None
|
||||
|
||||
|
||||
def test_a_corrupt_key_file_is_reported_rather_than_read_as_no_keys(tmp_path, caplog):
|
||||
"""Silently becoming "no keys configured" reads to the user as "my keys
|
||||
vanished", with the cause nowhere."""
|
||||
path = tmp_path / "inbound-keys.json"
|
||||
path.write_text("{not json", encoding="utf-8")
|
||||
|
||||
with caplog.at_level("ERROR"):
|
||||
store = KeyStore(str(path))
|
||||
|
||||
assert store.list_keys() == []
|
||||
assert "unreadable" in caplog.text
|
||||
|
||||
|
||||
def test_authentication_records_who_connected_and_from_where(store):
|
||||
issued = store.issue("Alice laptop")
|
||||
|
||||
store.authenticate(issued.secret, peer="10.0.0.5")
|
||||
|
||||
row = store.list_keys()[0]
|
||||
assert row["label"] == "Alice laptop"
|
||||
assert row["last_seen_peer"] == "10.0.0.5"
|
||||
assert row["last_seen_at"] > 0
|
||||
|
||||
|
||||
# ── Connection string ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_the_connection_string_round_trips(store):
|
||||
issued = store.issue("Alice")
|
||||
text = format_connection(host="192.168.0.110", port=7444, secret=issued.secret)
|
||||
|
||||
parsed = parse_connection(text)
|
||||
|
||||
assert parsed.host == "192.168.0.110"
|
||||
assert parsed.port == 7444
|
||||
assert parsed.secret == issued.secret
|
||||
assert parsed.endpoint == "192.168.0.110:7444"
|
||||
|
||||
|
||||
def test_an_ipv6_node_is_bracketed_for_grpc():
|
||||
"""gRPC's resolver reads an unbracketed IPv6 address as host:port and
|
||||
fails on the wrong half of it."""
|
||||
text = format_connection(host="fd00::1", port=7444, secret=KEY_PREFIX + "a" * 32)
|
||||
|
||||
assert parse_connection(text).endpoint == "[fd00::1]:7444"
|
||||
|
||||
|
||||
def test_the_secret_never_appears_in_the_loggable_form():
|
||||
connection = parse_connection(
|
||||
format_connection(host="10.0.0.2", port=7444, secret=KEY_PREFIX + "s" * 40)
|
||||
)
|
||||
|
||||
redacted = connection.redacted()
|
||||
|
||||
assert connection.secret not in redacted
|
||||
assert "10.0.0.2:7444" in redacted
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text, expected",
|
||||
[
|
||||
("192.168.0.110:7444", "without a key"),
|
||||
("ovnode://192.168.0.110:7444", "no key in it"),
|
||||
("https://192.168.0.110:7444", "connection string"),
|
||||
("ovnode://ovnode_short@10.0.0.1:7444", "not in the expected format"),
|
||||
("", "Paste the connection string"),
|
||||
],
|
||||
)
|
||||
def test_a_malformed_connection_string_says_what_is_wrong(text, expected):
|
||||
"""Every one of these otherwise surfaces as "cannot connect", which is the
|
||||
same thing a firewall, a wrong port and a dead node all say."""
|
||||
with pytest.raises(InvalidConnectionString) as excinfo:
|
||||
parse_connection(text)
|
||||
|
||||
assert expected in str(excinfo.value)
|
||||
|
||||
|
||||
def test_pasting_an_outbound_enrollment_token_says_so():
|
||||
"""The two credentials look alike and go in opposite directions; "invalid
|
||||
key" would send someone hunting for a typo that is not there."""
|
||||
with pytest.raises(InvalidConnectionString) as excinfo:
|
||||
parse_connection("ovnode://ovw_" + "a" * 40 + "@10.0.0.1:7444")
|
||||
|
||||
assert "other direction" in str(excinfo.value)
|
||||
Reference in New Issue
Block a user