Bound warm workers' host memory and account local RAM

A worker over FLOW_WORKER_MAX_RSS_MB (--worker-max-rss) when its node
returns is retired instead of kept warm, in the shared pool and in every
declared one; 0, the default, keeps them all. The engine now accounts its
own memory (FLOW_RAM / --ram, else the box or the cgroup limit less 1 GB),
so a node's `ram` declaration holds locally. A declared environment the
shared pool already has is that pool rather than a duplicate set of warm
workers.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

backend/fluksio/cli.py

@@ -262,6 +262,8 @@ CONCURRENCY_FLAGS = {
     "max_cascades": "FLOW_MAX_CASCADES",
     "max_runs": "FLOW_MAX_RUNS",
     "gpus": "FLOW_GPUS",
+    "ram": "FLOW_RAM",
+    "worker_max_rss": "FLOW_WORKER_MAX_RSS_MB",
 }
 
 
@@ -751,6 +753,20 @@ def _parser() -> argparse.ArgumentParser:
         metavar="N",
         help="GPUs a node may be given (default: /dev/nvidia* counted, FLOW_GPUS)",
     )
+    serve.add_argument(
+        "--ram",
+        type=_at_least(0),
+        default=None,
+        metavar="MB",
+        help="memory a node may be given (default: worked out, FLOW_RAM)",
+    )
+    serve.add_argument(
+        "--worker-max-rss",
+        type=_at_least(0),
+        default=None,
+        metavar="MB",
+        help="retire a warm worker holding more (default 0: never, FLOW_WORKER_MAX_RSS_MB)",
+    )
     serve.set_defaults(func=cmd_serve)
 
     enroll = subparsers.add_parser("enroll", help="pair this instance with a portal")

backend/fluksio/core/config.py

@@ -142,6 +142,16 @@ class Settings(BaseSettings):
     # NVIDIA device nodes, the way 0 cores works the core count out; anything
     # else — another vendor, or keeping a card back — is said outright.
     FLOW_GPUS: int = 0
+    # Megabytes nodes that declare `ram` may be given. 0 works it out: the
+    # machine's memory, or the container's limit where that is lower, less a
+    # gigabyte for the engine. Like cores, a node declaring nothing is not
+    # counted against it.
+    FLOW_RAM: int = 0
+    # A warm worker is retired rather than kept when its node returns with the
+    # process holding more than this many megabytes resident. A library that
+    # grew it then gives the memory back, and the next call pays its imports
+    # again. 0, the default, keeps every worker warm.
+    FLOW_WORKER_MAX_RSS_MB: int = 0
     # How long the engine's own metrics, events and run records are kept.
     #: The largest body `PUT /artifacts` will take, in bytes. A checkpoint or a
     #: video segment is a legitimate artifact, so this is generous rather than

backend/fluksio/flow/resources.py

@@ -31,6 +31,7 @@ import re
 import threading
 from collections.abc import Callable
 from dataclasses import dataclass
+from pathlib import Path
 
 from sqlmodel import Session
 
@@ -58,6 +59,12 @@ THREAD_VARS = (
 #: in that process, and a node is not allowed to take the last of it.
 ENGINE_RESERVE = 2
 
+#: The same for memory, in megabytes.
+ENGINE_RAM_RESERVE_MB = 1024
+
+#: Where a container finds its own memory limit (cgroup v2).
+CGROUP_MEMORY_MAX = Path("/sys/fs/cgroup/memory.max")
+
 
 def machine_cpus() -> int:
     """Cores a node may be given here, when nobody said."""
@@ -76,6 +83,27 @@ def machine_gpus() -> int:
     return len(glob.glob("/dev/nvidia[0-9]*"))
 
 
+def machine_ram_mb() -> int | None:
+    """Megabytes a node may be given here, when nobody said.
+
+    The box's memory, or the container's limit where that is lower — the
+    host's total is memory the engine would be killed for using. None where
+    neither can be asked, which leaves memory unaccounted rather than zero.
+    """
+    try:
+        total = os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") // 2**20
+    except (AttributeError, ValueError, OSError):
+        return None
+    try:
+        # "max" when the container has no limit of its own.
+        limit = CGROUP_MEMORY_MAX.read_text().strip()
+    except OSError:
+        limit = ""
+    if limit.isdigit():
+        total = min(total, int(limit) // 2**20)
+    return max(1, total - ENGINE_RAM_RESERVE_MB)
+
+
 @dataclass(frozen=True, slots=True)
 class Allocation:
     """What one execution was actually given."""

backend/fluksio/flow/workers.py

@@ -153,6 +153,17 @@ class _Worker:
     def alive(self) -> bool:
         return self.proc.poll() is None
 
+    def rss_mb(self) -> int:
+        """Resident memory in megabytes, or 0 where it cannot be asked."""
+        # ponytail: /proc is Linux-only; elsewhere every worker reads as small
+        # and the ceiling never retires one.
+        try:
+            with open(f"/proc/{self.proc.pid}/statm") as statm:
+                pages = int(statm.read().split()[1])
+        except (OSError, ValueError, IndexError):
+            return 0
+        return pages * os.sysconf("SC_PAGE_SIZE") // 2**20
+
     def send(self, request: dict[str, Any]) -> None:
         assert self.proc.stdin is not None
         # Node inputs may be keyed by something other than a string, which
@@ -221,10 +232,16 @@ class PythonWorkerPool:
         size: int = 4,
         events: EventBus | None = None,
         env: dict[str, str] | None = None,
+        max_rss_mb: int = 0,
     ) -> None:
         self.python = python
         self.size = size
         self.events = events
+        # A worker holding more than this when its call returns is retired
+        # rather than kept warm; 0 keeps them all. A library like JAX keeps what
+        # it allocated for as long as the process lives, and warm is the point
+        # of a pool, so nothing else would ever give that memory back.
+        self.max_rss_mb = max_rss_mb
         # 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 {}
@@ -271,8 +288,10 @@ class PythonWorkerPool:
 
         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.
+        An environment this pool's workers already have is this pool, rather
+        than a second set of warm workers holding the same imports.
         """
-        if not extra:
+        if extra.items() <= self.env.items():
             return self
         key = frozenset(extra.items())
         with self._lock:
@@ -283,6 +302,7 @@ class PythonWorkerPool:
                     size=self.size,
                     events=self.events,
                     env={**self.env, **extra},
+                    max_rss_mb=self.max_rss_mb,
                 )
                 child._generation = self._generation
                 child.start()
@@ -379,6 +399,7 @@ class PythonWorkerPool:
             worker.alive()
             and worker.generation == self._generation
             and not worker.cancelled
+            and not (self.max_rss_mb and worker.rss_mb() > self.max_rss_mb)
         )
         if reusable:
             self._idle.put(worker)

backend/fluksio/main.py

@@ -44,6 +44,7 @@ from fluksio.flow.resources import (
     ResourceAccountant,
     fair_share_env,
     machine_gpus,
+    machine_ram_mb,
 )
 from fluksio.flow.runs import RUN_STATE_TTL, RunService, sweep_artifacts
 from fluksio.flow.secrets import init_secrets
@@ -256,9 +257,11 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
         volatile = VolatileStore(_volatile_root(), settings.ARTIFACT_VOLATILE_BYTES)
         artifacts.volatile = volatile
         app.state.artifact_store = artifacts
-        # Zero means "work it out" for both, as it always has for cores.
+        # Zero means "work it out" for cards and memory, as it always has for cores.
         accountant = ResourceAccountant(
-            cpus=settings.FLOW_CPUS, gpus=settings.FLOW_GPUS or machine_gpus()
+            cpus=settings.FLOW_CPUS,
+            gpus=settings.FLOW_GPUS or machine_gpus(),
+            ram_mb=settings.FLOW_RAM or machine_ram_mb(),
         )
         app.state.resources = accountant
         # Every machine a node could run on: this one, and whatever attaches.
@@ -270,6 +273,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
             python=modules.venv_python(),
             size=settings.FLOW_MAX_WORKERS,
             events=event_bus,
+            max_rss_mb=settings.FLOW_WORKER_MAX_RSS_MB,
             # A worker in this container writes to the store directly; a remote one
             # is given a URL instead. Node code calls the same two functions.
             env={
@@ -332,10 +336,12 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
         # Said out loud because it is the only way to tell that a settings file was
         # read at all — the numbers are what someone raising them is looking for.
         logger.info(
-            "Engine limits: workers=%s cascades=%s runs=%s",
+            "Engine limits: workers=%s cascades=%s runs=%s ram_mb=%s worker_max_rss_mb=%s",
             settings.FLOW_MAX_WORKERS,
             execution.max_cascades,
             run_service.parallel,
+            accountant.ram_mb,
+            settings.FLOW_WORKER_MAX_RSS_MB,
         )
         watchdog = LoopWatchdog(event_bus)
         app.state.watchdog = watchdog

backend/tests/flow/test_resources.py

@@ -19,6 +19,7 @@ from fluksio.flow.resources import (
     derive_env,
     fair_share_env,
     machine_gpus,
+    machine_ram_mb,
 )
 from fluksio.flow.schemas import NodeDef, Resources
 
@@ -44,6 +45,20 @@ def test_the_cards_are_counted_from_the_devices_the_driver_makes(monkeypatch):
     assert machine_gpus() == 0
 
 
+def test_memory_is_the_containers_limit_where_that_is_lower(monkeypatch, tmp_path):
+    """The host's total is memory a container would be killed for using."""
+    from fluksio.flow import resources
+
+    box = os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") // 2**20
+    limit = tmp_path / "memory.max"
+    monkeypatch.setattr(resources, "CGROUP_MEMORY_MAX", limit)
+
+    limit.write_text("max\n")
+    assert machine_ram_mb() == box - resources.ENGINE_RAM_RESERVE_MB
+    limit.write_text(f"{2048 * 2**20}\n")
+    assert machine_ram_mb() == 2048 - resources.ENGINE_RAM_RESERVE_MB
+
+
 def test_what_is_free_is_what_was_handed_out():
     accountant = ResourceAccountant(cpus=4, gpus=0)
     held = accountant.try_take(cpus=3, gpus=0)

backend/tests/flow/test_workers.py

@@ -659,6 +659,9 @@ def test_a_declared_environment_gets_a_pool_of_its_own(pool):
     assert pool.for_env({"OMP_NUM_THREADS": "3"}) is not child
     # Nothing to derive means the shared pool, with no second process anywhere.
     assert pool.for_env({}) is pool
+    # Nor does deriving what the shared pool's workers were started with.
+    pool.env["OMP_NUM_THREADS"] = "4"
+    assert pool.for_env({"OMP_NUM_THREADS": "4"}) is pool
 
 
 def test_a_child_worker_is_started_with_what_it_was_given(pool):
@@ -698,6 +701,23 @@ def test_retiring_the_cards_leaves_the_other_pools_warm(pool):
     assert threads._generation == before[1]
 
 
+def test_a_worker_over_the_memory_ceiling_is_not_kept_warm(pool):
+    """A library keeps what it allocated, so the process is what gives it back."""
+    code = "import os\n\n\ndef process():\n    return {'pid': os.getpid()}\n"
+    # Any interpreter is past a megabyte, so every call retires its worker.
+    tight = PythonWorkerPool(python=sys.executable, size=1, max_rss_mb=1)
+    tight.start()
+    try:
+        first = tight.run("demo", "pid", code, {}, "demo.pid", timeout=10)
+        assert tight.run("demo", "pid", code, {}, "demo.pid", timeout=10) != first
+        # The pools a declaration gets are held to the same ceiling.
+        assert tight.for_env({"OMP_NUM_THREADS": "2"}).max_rss_mb == 1
+    finally:
+        tight.stop()
+    # Without one, the same slot answers twice from the same process.
+    assert run(pool, code, node="pid") == run(pool, code, node="pid")
+
+
 def test_cancelling_reaches_a_node_running_in_a_child(pool):
     child = pool.for_env({"OMP_NUM_THREADS": "2"})
     started = threading.Event()