Let a node say how much of the machine it takes
Five concurrent training nodes, each sizing its thread pool to every core,
left the engine's own event loop unscheduled: the API stopped answering
within 10 s and every client died. The same shape on a GPU deadlocked a run
for 21 minutes at 0% utilisation with nothing failing and nothing to read --
it just sat in `running`.
@node(resources={"cpus": 2}) is the declaration. The engine holds that much
for the length of the execution, so more of them than the machine has room
for wait their turn rather than oversubscribing it, and a `gpus` node holds
its card exclusively. FLOW_CPUS defaults to every core but two, and those two
are what keeps the engine answering.
Because a thread cap is read when the process imports the library, a warm
worker cannot be told a different one -- so an environment gets a pool of its
own and nodes deriving the same one share it, rather than paying a cold start
per call on exactly the nodes whose imports are slowest. XLA_FLAGS is never
derived: it is a composed, version-dependent string, so it travels in
resources.env where it is visible.
A node that declares nothing is not accounted for and behaves as it always
did -- it just gets FLOW_CPUS/FLOW_MAX_WORKERS as a thread cap, which is the
half of this that fixes the reported incident without anybody declaring
anything. An operator who set OMP_NUM_THREADS themselves still wins.
Resources are claimed strictly before a worker slot, so the two blocking
waits cannot deadlock. A node queued for them publishes node_queued and shows
on GET /workers/resources, because waiting and hanging looked identical.
Accounted, not enforced: no cgroups, no rlimits. Scheduling across machines,
flavours and enforcement are the next steps.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -228,6 +228,10 @@ class PythonWorkerPool:
|
||||
# Put into every worker's environment past the denylist — where the
|
||||
# artifact store is, which node code needs and cannot guess.
|
||||
self.env = env or {}
|
||||
# Pools for nodes that declared resources, keyed by the environment
|
||||
# their allocation derived. Empty on all but the one pool the engine
|
||||
# builds; see ``for_env``.
|
||||
self._children: dict[frozenset[tuple[str, str]], PythonWorkerPool] = {}
|
||||
self._idle: queue.Queue[_Worker | None] = queue.Queue()
|
||||
# Keyed by (run, node): a sweep has the same node executing in several
|
||||
# runs at once, and cancelling one of them must not kill the others.
|
||||
@@ -255,9 +259,41 @@ class PythonWorkerPool:
|
||||
for _ in range(self.size):
|
||||
self._idle.put(None)
|
||||
|
||||
def for_env(self, extra: dict[str, str]) -> PythonWorkerPool:
|
||||
"""The pool whose workers were started with this environment.
|
||||
|
||||
A thread cap or a GPU assignment is read once, when the process
|
||||
imports the library, so a worker that is already warm cannot be told a
|
||||
different one — and retiring a worker per call would give back the cold
|
||||
starts ``_warm`` exists to avoid, on exactly the nodes whose imports
|
||||
are slowest. So an environment gets a pool of its own instead, and
|
||||
nodes deriving the same one share it.
|
||||
|
||||
How many of them may run at once is the accountant's business, not this
|
||||
pool's: the slots here only bound one environment's own concurrency.
|
||||
"""
|
||||
if not extra:
|
||||
return self
|
||||
key = frozenset(extra.items())
|
||||
with self._lock:
|
||||
child = self._children.get(key)
|
||||
if child is None:
|
||||
child = PythonWorkerPool(
|
||||
self.python,
|
||||
size=self.size,
|
||||
events=self.events,
|
||||
env={**self.env, **extra},
|
||||
)
|
||||
child._generation = self._generation
|
||||
child.start()
|
||||
self._children[key] = child
|
||||
return child
|
||||
|
||||
def stop(self) -> None:
|
||||
self._stopped = True
|
||||
self._generation += 1
|
||||
for child in list(self._children.values()):
|
||||
child.stop()
|
||||
for worker in list(self._running.values()):
|
||||
worker.kill()
|
||||
for slot in self._drain():
|
||||
@@ -281,6 +317,9 @@ class PythonWorkerPool:
|
||||
if slot is not None:
|
||||
slot.kill()
|
||||
self._idle.put(None)
|
||||
children = list(self._children.values())
|
||||
for child in children:
|
||||
child.respawn_all()
|
||||
|
||||
def _drain(self) -> list[_Worker | None]:
|
||||
slots = []
|
||||
@@ -575,11 +614,14 @@ class PythonWorkerPool:
|
||||
"""Stop a node that is running now. False when there was nothing to stop."""
|
||||
with self._lock:
|
||||
worker = self._running.get((run_id, node_id))
|
||||
if worker is None:
|
||||
return False
|
||||
worker.cancelled = True
|
||||
worker.kill()
|
||||
return True
|
||||
children = list(self._children.values())
|
||||
if worker is not None:
|
||||
worker.cancelled = True
|
||||
worker.kill()
|
||||
return True
|
||||
# A node with declared resources runs in a pool of its own, and it is
|
||||
# the same node to whoever pressed cancel.
|
||||
return any(child.cancel(node_id, run_id) for child in children)
|
||||
|
||||
def cancel_run(self, run_id: str) -> int:
|
||||
"""Stop every node this run has in a worker right now."""
|
||||
@@ -592,7 +634,8 @@ class PythonWorkerPool:
|
||||
for worker in workers:
|
||||
worker.cancelled = True
|
||||
worker.kill()
|
||||
return len(workers)
|
||||
children = list(self._children.values())
|
||||
return len(workers) + sum(child.cancel_run(run_id) for child in children)
|
||||
|
||||
def _publish(self, event: dict[str, Any]) -> None:
|
||||
if self.events is not None:
|
||||
|
||||
Reference in New Issue
Block a user