Let a worker say what machine it is

A worker reported its labels and nothing about the machine behind them, so the
engine could route a node to a GPU box but not tell whether that box had a GPU
free. Inventory — cores, GPUs, memory — now arrives with the hello frame, and
the run frame carries back what the engine allocated for that call.

Which is protocol 2 on both ends. GPUs are never probed: asking a vendor tool
would make the one dependency two, so a GPU is what the batch job says it was
given or what --gpus says. A worker that reports nothing still attaches and is
scheduled by its label alone.

Two things a job scheduler needs: --max-idle stops a worker started for one job
rather than letting it hold its allocation to the walltime, and a refusal is now
fatal instead of a reconnect loop that reads as a hang in a job's log.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01A6HeySA27EkGANZN95QySW

backend/fluksio/api/routes/workers.py

@@ -21,7 +21,12 @@ from pydantic import BaseModel, Field
 
 from fluksio.api.deps import get_current_active_superuser, get_current_user
 from fluksio.core import security
-from fluksio.flow.remote import PROTOCOL, RemoteWorker, RemoteWorkerHub
+from fluksio.flow.remote import (
+    PROTOCOL,
+    RemoteWorker,
+    RemoteWorkerHub,
+    WorkerInventory,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -40,6 +45,9 @@ class WorkerInfo(BaseModel):
     last_seen: float = 0.0
     python: str = ""
     venv_digest: str = ""
+    cpus: int = 1
+    gpus: int = 0
+    ram_mb: int | None = None
 
 
 class TokenRequest(BaseModel):
@@ -74,6 +82,9 @@ def read_workers(request: Request) -> Any:
             last_seen=worker.last_seen,
             python=str(worker.info.get("python") or ""),
             venv_digest=str(worker.info.get("venv_digest") or ""),
+            cpus=worker.inventory.cpus,
+            gpus=worker.inventory.gpus,
+            ram_mb=worker.inventory.ram_mb,
         )
         for worker in _hub(request.app).workers()
     ]
@@ -163,6 +174,7 @@ async def attach(websocket: WebSocket, token: str = "") -> None:
             "python": hello.get("python"),
             "venv_digest": hello.get("venv_digest"),
         },
+        inventory=WorkerInventory.from_hello(hello.get("inventory")),
     )
     hub.attach(worker)
     await websocket.send_json({"op": "welcome", "protocol": PROTOCOL, "name": name})

backend/fluksio/flow/remote.py

@@ -28,6 +28,7 @@ import sys
 import threading
 import time
 from collections.abc import Callable
+from dataclasses import dataclass
 from typing import Any
 
 from fluksio.flow.nodes.base import NodeOutputError
@@ -44,13 +45,50 @@ SEND_TIMEOUT_S = 30.0
 SILENCE_S = 90.0
 #: Protocol version this engine speaks. A worker announcing anything else is
 #: refused rather than half-understood.
-PROTOCOL = 1
+PROTOCOL = 2
 
 
 class NoWorker(RemoteError):
     """Nothing is attached that carries the label this node asked for."""
 
 
+def _whole(value: Any, default: int) -> int:
+    """An int from whatever a worker sent, or the default when it was junk."""
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return default
+
+
+@dataclass(frozen=True)
+class WorkerInventory:
+    """What an attached worker says it holds.
+
+    The defaults are a machine that runs one thing at a time and has no GPU,
+    so a worker reporting nothing is scheduled by its label alone — which is
+    how every worker was scheduled before any of them reported anything.
+    """
+
+    cpus: int = 1
+    gpus: int = 0
+    ram_mb: int | None = None
+
+    @classmethod
+    def from_hello(cls, payload: Any) -> WorkerInventory:
+        """Read what is understood and ignore the rest.
+
+        Unknown members are deliberately not an error: a worker that learns to
+        report something new must not need this engine taught about it first.
+        """
+        if not isinstance(payload, dict):
+            return cls()
+        return cls(
+            cpus=max(1, _whole(payload.get("cpus"), 1)),
+            gpus=max(0, _whole(payload.get("gpus"), 0)),
+            ram_mb=_whole(payload.get("ram_mb"), 0) or None,
+        )
+
+
 class RemoteWorker:
     """One attached worker, and the calls it has in flight."""
 
@@ -62,10 +100,12 @@ class RemoteWorker:
         loop: asyncio.AbstractEventLoop,
         max_parallel: int = 1,
         info: dict[str, Any] | None = None,
+        inventory: WorkerInventory | None = None,
     ) -> None:
         self.name = name
         self.labels = set(labels)
         self.info = info or {}
+        self.inventory = inventory or WorkerInventory()
         self.attached_at = time.time()
         self.last_seen = time.time()
         self._send = send
@@ -281,21 +321,28 @@ class RemoteWorkerHub:
         timeout: float,
         run_id: str = "",
         on_event: Callable[[dict[str, Any]], None] | None = None,
+        env: dict[str, str] | None = None,
     ) -> Any:
         worker = self.pick(label)
         if worker is None:
             raise NoWorker(f"no worker labelled '{label}' is attached")
+        payload = {
+            "op": "run",
+            "call_id": f"{run_id}:{node_id}" if run_id else node_id,
+            "flow": flow,
+            "node": node,
+            "source": source,
+            "kwargs": kwargs,
+            "run": {"id": run_id} if run_id else None,
+            "timeout": timeout,
+        }
+        if env:
+            # What this call was allocated: the thread caps and the devices the
+            # node may see. The worker starts a process per call, so it applies
+            # them where a library still reads them — before the import.
+            payload["env"] = env
         response = worker.request(
-            {
-                "op": "run",
-                "call_id": f"{run_id}:{node_id}" if run_id else node_id,
-                "flow": flow,
-                "node": node,
-                "source": source,
-                "kwargs": kwargs,
-                "run": {"id": run_id} if run_id else None,
-                "timeout": timeout,
-            },
+            payload,
             timeout=timeout,
             on_event=on_event,
         )
@@ -357,6 +404,7 @@ class RemoteWorkerHub:
         run_id: str = "",
         on_event: Callable[[dict[str, Any]], None] | None = None,
         fallback: Callable[..., Any] | None = None,
+        env: dict[str, str] | None = None,
     ) -> Callable[..., Any]:
         """What a node with a device runs instead of its own function.
 
@@ -377,6 +425,7 @@ class RemoteWorkerHub:
                 timeout,
                 run_id=run_id,
                 on_event=on_event,
+                env=env,
             )
 
         return call

backend/tests/flow/test_remote.py

@@ -15,7 +15,12 @@ from collections.abc import Iterator
 import pytest
 
 from fluksio.flow import remote
-from fluksio.flow.remote import NoWorker, RemoteWorker, RemoteWorkerHub
+from fluksio.flow.remote import (
+    NoWorker,
+    RemoteWorker,
+    RemoteWorkerHub,
+    WorkerInventory,
+)
 from fluksio.flow.workers import NodeTimeout, RemoteError
 
 
@@ -295,3 +300,67 @@ def test_with_no_timeout_a_beating_worker_is_left_to_finish(loop, monkeypatch):
 
     thread.join(timeout=5)
     assert result["value"] == {"done": True}
+
+
+def test_an_allocation_rides_along_with_the_call(loop):
+    """Thread caps and devices reach the worker on the frame that needs them.
+
+    The worker starts a process per call, so this is the only moment it can
+    apply them: a library reads them when it is imported and never again.
+    """
+    hub = RemoteWorkerHub()
+    worker, socket = attach(hub, loop)
+
+    thread = call_in_thread(
+        lambda: hub.run(
+            "gpu",
+            "flow",
+            "node",
+            "src",
+            {},
+            "flow.node",
+            timeout=5,
+            env={"OMP_NUM_THREADS": "4", "CUDA_VISIBLE_DEVICES": "0"},
+        )
+    )
+    assert socket.arrived.wait(5)
+    assert socket.sent[0]["env"] == {
+        "OMP_NUM_THREADS": "4",
+        "CUDA_VISIBLE_DEVICES": "0",
+    }
+
+    worker.deliver({"call_id": socket.sent[0]["call_id"], "ok": True, "result": None})
+    thread.join(timeout=5)
+
+
+def test_a_call_with_nothing_allocated_carries_no_env(loop):
+    hub = RemoteWorkerHub()
+    worker, socket = attach(hub, loop)
+
+    thread = call_in_thread(
+        lambda: hub.run("gpu", "flow", "node", "src", {}, "flow.node", timeout=5)
+    )
+    assert socket.arrived.wait(5)
+    assert "env" not in socket.sent[0]
+
+    worker.deliver({"call_id": socket.sent[0]["call_id"], "ok": True, "result": None})
+    thread.join(timeout=5)
+
+
+def test_a_worker_that_reports_nothing_is_the_machine_it_used_to_be():
+    """Inventory is read leniently: absent is a default, junk is a default.
+
+    A worker saying nothing has to keep being scheduled by its label, and one
+    that learns to report something new must not need this engine taught about
+    it first.
+    """
+    assert WorkerInventory.from_hello({}) == WorkerInventory(cpus=1, gpus=0)
+    assert WorkerInventory.from_hello(None) == WorkerInventory()
+
+    read = WorkerInventory.from_hello(
+        {"cpus": "8", "gpus": 2, "ram_mb": 64000, "tpus": 4}
+    )
+    assert (read.cpus, read.gpus, read.ram_mb) == (8, 2, 64000)
+
+    junk = WorkerInventory.from_hello({"cpus": "many", "gpus": None, "ram_mb": 0})
+    assert (junk.cpus, junk.gpus, junk.ram_mb) == (1, 0, None)