From dafdcb76c4b4d04bc3f217c46228a5bcb3f63c40 Mon Sep 17 00:00:00 2001 From: stroblme Date: Sun, 27 Sep 2026 16:11:13 +0200 Subject: [PATCH] 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 --- backend/fluksio/cli.py | 16 ++++++++++++++++ backend/fluksio/core/config.py | 10 ++++++++++ backend/fluksio/flow/resources.py | 28 ++++++++++++++++++++++++++++ backend/fluksio/flow/workers.py | 23 ++++++++++++++++++++++- backend/fluksio/main.py | 12 +++++++++--- backend/tests/flow/test_resources.py | 15 +++++++++++++++ backend/tests/flow/test_workers.py | 20 ++++++++++++++++++++ docs/code/cli.md | 2 ++ docs/getting-started/data-science.md | 7 +++++++ docs/reference/configuration.md | 18 ++++++++++++++---- 10 files changed, 143 insertions(+), 8 deletions(-) diff --git a/backend/fluksio/cli.py b/backend/fluksio/cli.py index c0466de..b187f48 100644 --- a/backend/fluksio/cli.py +++ b/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") diff --git a/backend/fluksio/core/config.py b/backend/fluksio/core/config.py index 3dd2e6b..d58cbe5 100644 --- a/backend/fluksio/core/config.py +++ b/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 diff --git a/backend/fluksio/flow/resources.py b/backend/fluksio/flow/resources.py index 27cd6ea..be51011 100644 --- a/backend/fluksio/flow/resources.py +++ b/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.""" diff --git a/backend/fluksio/flow/workers.py b/backend/fluksio/flow/workers.py index 4abf517..7cb4b78 100644 --- a/backend/fluksio/flow/workers.py +++ b/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) diff --git a/backend/fluksio/main.py b/backend/fluksio/main.py index fef3049..fccd320 100644 --- a/backend/fluksio/main.py +++ b/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 diff --git a/backend/tests/flow/test_resources.py b/backend/tests/flow/test_resources.py index 5276028..f4ee54e 100644 --- a/backend/tests/flow/test_resources.py +++ b/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) diff --git a/backend/tests/flow/test_workers.py b/backend/tests/flow/test_workers.py index 1e9798b..0a4c47e 100644 --- a/backend/tests/flow/test_workers.py +++ b/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() diff --git a/docs/code/cli.md b/docs/code/cli.md index 70ae500..8b11e47 100644 --- a/docs/code/cli.md +++ b/docs/code/cli.md @@ -79,6 +79,8 @@ Another instance's Fluksio on the port is named, and the move happens as usual. | `--max-cascades N` | 4 | cascades in flight at once (`FLOW_MAX_CASCADES`) | | `--max-workers N` | 4 | python worker processes (`FLOW_MAX_WORKERS`) | | `--gpus N` | counted | GPUs on this machine a node may be given (`FLOW_GPUS`) | +| `--ram MB` | worked out | memory on this machine a node may be given (`FLOW_RAM`) | +| `--worker-max-rss MB` | 0 (never) | retire a warm worker holding more when its node returns (`FLOW_WORKER_MAX_RSS_MB`) | Cards are counted from NVIDIA's device nodes (`/dev/nvidia0`, `/dev/nvidia1`, …), which asks no vendor tool and so keeps the one dependency from becoming two. diff --git a/docs/getting-started/data-science.md b/docs/getting-started/data-science.md index 0086c3c..89caec9 100644 --- a/docs/getting-started/data-science.md +++ b/docs/getting-started/data-science.md @@ -315,6 +315,13 @@ given `CUDA_VISIBLE_DEVICES`, and leaves a warm worker holding whatever it preallocated. Declaring the card is what makes it a worker the engine knows to retire, and what stops two such nodes running at once in the first place. +Host memory is held the same way, and there the declaration is not needed. A +worker that fitted a model keeps what it allocated, and on a laptop a few warm +ones are the machine. `fluksio serve --worker-max-rss 1024` retires any worker +holding more than that when its node returns, whichever pool it is in. A +node's `ram` is counted against this machine's memory, so declaring it is what +stops three large fits from starting at once. + Declaring nothing is the default and is right for most nodes: a poll, a threshold, a message on its way somewhere. Those share the engine's worker pool and are given a fair share of `FLOW_CPUS` as a thread cap, which is what stops diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index 2193570..bacbf23 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -134,6 +134,8 @@ warning into a refusal to start. | `FLOW_NODE_TIMEOUT` | `0` | seconds a node may be silent, unless it sets its own; 0 is no limit | | `FLOW_CPUS` | `0` | cores nodes that declare `resources` may be given; 0 works it out as every core but two, which are what keeps the engine answering while the machine is busy | | `FLOW_GPUS` | `0` | GPUs on this machine, each held by one node at a time. 0 counts NVIDIA's device nodes; say a number for anything they miss | +| `FLOW_RAM` | `0` | megabytes nodes that declare `ram` may be given; 0 works it out as the machine's memory, or the container's limit where that is lower, less a gigabyte for the engine | +| `FLOW_WORKER_MAX_RSS_MB` | `0` | a warm worker holding more resident memory than this when its node returns is retired rather than kept; 0 keeps every worker warm | | `OBS_RETENTION_DAYS` | `30` | how long metrics, events and run records are kept | | `ARTIFACT_GC_INTERVAL_S` | `3600` | how often artifact bytes nothing refers to are swept away; 0 never sweeps | | `ARTIFACT_GC_GRACE_S` | `3600` | how long a freshly written artifact is spared, whatever refers to it | @@ -142,10 +144,10 @@ warning into a refusal to start. | `ARTIFACT_VOLATILE_BYTES` | `50331648` | how much that ring holds before the oldest frames fall out; 0 turns it off and volatile saves land in the store | The three concurrency limits are also flags on `fluksio serve` -(`--max-workers`, `--max-cascades`, `--max-runs`), as is the card count, -`--gpus`. The flags +(`--max-workers`, `--max-cascades`, `--max-runs`), as are the card count and +memory, `--gpus` and `--ram`, and the worker ceiling, `--worker-max-rss`. The flags outrank the file, and the engine says which numbers it started with in its -first lines. Each pool size must be at least 1 and the card count at least 0: +first lines. Each pool size must be at least 1 and the others at least 0: a number below that is refused as a flag error naming it, rather than read as the default. Leave one empty (or unset) to get the default. @@ -170,7 +172,15 @@ several at once cannot each size themselves to the whole machine. Setting `OMP_NUM_THREADS` (or any of its siblings) on the engine yourself overrides that default. -These two are this machine's figures. An attached worker reports its own when +Workers are kept warm between calls, and a library like JAX keeps what it +allocated for as long as its process lives — so a worker that fitted a model +can sit on gigabytes after its run is over, on the shared pool as much as on a +declared one. `FLOW_WORKER_MAX_RSS_MB` retires such a worker when its node +returns, and the next call pays its imports again. That bounds idle memory +only; what a node holds while it runs is accounted, and only if it declared +`ram`. + +`FLOW_CPUS`, `FLOW_GPUS` and `FLOW_RAM` are this machine's figures. An attached worker reports its own when it dials in, and a node goes to whichever machine can grant what it asked for, so a GPU on a worker needs no `FLOW_GPUS` here. What every machine has free, and which nodes are queued, is `GET /api/v1/workers/resources` and the Workers