diff --git a/.github/workflows/lre.yaml b/.github/workflows/lre.yaml index 73bb9b276a3..2c29e94f5e0 100644 --- a/.github/workflows/lre.yaml +++ b/.github/workflows/lre.yaml @@ -20,7 +20,41 @@ concurrency: cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} jobs: + # Compute whether this diff can affect the LRE example builds. These lanes + # build @local-remote-execution//examples, so they are keyed off the LRE + # toolchain inputs. Pushes to main always run the full matrix (backstop). + changes: + name: Compute affected lanes + runs-on: ubuntu-24.04 + outputs: + lre: ${{ steps.filter.outputs.lre }} + steps: + - name: Checkout + uses: >- # v6.0.2 + actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd + + - name: Filter LRE-affecting paths + uses: >- # v3.0.2 + dorny/paths-filter@de90cc6fb38fc0963ad72b210f1f284cd68cea36 + id: filter + with: + filters: | + lre: + - 'local-remote-execution/**' + - 'MODULE.bazel' + - 'MODULE.bazel.lock' + - '*.bazelrc' + - 'flake.nix' + - 'flake.lock' + - 'tools/**' + - '.github/workflows/lre.yaml' + - '.github/actions/**' + local: + needs: [changes] + if: >- + github.ref == 'refs/heads/main' + || needs.changes.outputs.lre == 'true' strategy: fail-fast: false matrix: diff --git a/.github/workflows/nix.yaml b/.github/workflows/nix.yaml index 1dd19e3265c..8d9afd866ab 100644 --- a/.github/workflows/nix.yaml +++ b/.github/workflows/nix.yaml @@ -157,8 +157,45 @@ jobs: nativelink_attic_token: ${{ secrets.NATIVELINK_ATTIC_TOKEN }} nix_name: .#${{ matrix.test-name }}-with-nativelink-test + # Compute which lanes this diff can affect. The rbe-* lanes validate + # "does this toolchain still build through NativeLink", so they are keyed + # off the toolchain/LRE configuration inputs they consume. Pushes to main + # always run the full matrix (backstop), so no change can merge without + # the complete matrix running against it. + changes: + name: Compute affected lanes + runs-on: ubuntu-24.04 + outputs: + toolchain: ${{ steps.filter.outputs.toolchain }} + steps: + - name: Checkout + uses: >- # v6.0.2 + actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd + + - name: Filter toolchain-affecting paths + uses: >- # v3.0.2 + dorny/paths-filter@de90cc6fb38fc0963ad72b210f1f284cd68cea36 + id: filter + with: + filters: | + toolchain: + - 'toolchain-examples/**' + - 'local-remote-execution/**' + - 'MODULE.bazel' + - 'MODULE.bazel.lock' + - '*.bazelrc' + - 'flake.nix' + - 'flake.lock' + - 'tools/**' + - '.github/workflows/nix.yaml' + - '.github/actions/**' + # rbe-toolchain is slow because it has many subcommands, so split it out generate-rbe-commands: + needs: [changes] + if: >- + github.ref == 'refs/heads/main' + || needs.changes.outputs.toolchain == 'true' runs-on: ubuntu-24.04 outputs: matrix: ${{ steps.set-matrix.outputs.matrix }} diff --git a/llms-full.txt b/llms-full.txt index aa28b6a02c2..d2abea0b8e7 100644 --- a/llms-full.txt +++ b/llms-full.txt @@ -6306,16 +6306,15 @@ alive between them. ## Two things to know before you start [#two-things-to-know-before-you-start] -**Almost nothing needs configuring in NativeLink.** No field anywhere in -[the configuration model](/configuration) turns this on, tunes the pool, or -selects a protocol; `LocalWorkerConfig` has no persistent-worker section at -all, and every pool parameter is a compiled-in constant. The feature is -driven by two platform properties on each action. The one thing NativeLink -does need from you is on the scheduler: those two keys must be listed in -`supported_platform_properties`, because the scheduler rejects an action that -carries a key it does not know. If you have seen a `persistent_workers` -config block in older documentation, it did not correspond to anything in the -binary. +**Little needs configuring in NativeLink.** The pool is on by default and +the protocol is chosen by two platform properties on each action. The +worker's `persistent_workers` block (v1.7.3 and later) sizes it: how many +processes per key, how long an idle one lives, how many requests one serves, +how long an action at the cap waits for a process to come back, and +`enabled: false` to run every such action one-shot. +The one thing NativeLink does need from you is on the scheduler: those two +keys must be listed in `supported_platform_properties`, because a scheduler +older than v1.6.6 rejects an action that carries a key it does not know. **The module is licensed separately.** Most of the repository is `FSL-1.1-Apache-2.0`. Every file under `nativelink-worker/src/persistent_worker/` @@ -6435,19 +6434,23 @@ while everything still appears to work. ## The isolation you give up [#the-isolation-you-give-up] -This is the part to read twice. NativeLink applies real isolation to ordinary -actions and **applies none of it to persistent workers**. The persistent path -returns before any of it runs. +This is the part to read twice. A persistent worker process outlives the +actions it serves, so some of the isolation an ordinary action gets cannot +apply to it, and some of it did not apply until v1.7.3. -**No namespace sandboxing.** The PID, user, IPC, UTS and mount namespaces that -ordinary actions get are not applied. `use_namespaces` and `use_mount_namespace` -have no effect on a persistent worker. +**Process namespaces, but not the mount namespace.** With `use_namespaces` +the worker process gets its own PID, user, IPC and UTS namespaces, as an +ordinary action does, so its children are reaped with it. It never gets a +mount namespace or the private `/tmp` (`isolate_tmp`): both are per action, +and the process is not. Persistent workers see the host's `/tmp` and the +whole of the worker's filesystem, whatever `use_mount_namespace` says. -**No environment control.** Ordinary actions get a cleared environment plus -exactly what the action and config specify. A persistent worker process -inherits the NativeLink worker process's own environment, and -`environment_variables` on the command, `additional_environment` on the config, -and the side-channel environment sources are all ignored. +**The action's environment, and only that.** The process is started with a +cleared environment plus the command's `environment_variables`, the way an +ordinary action is, and that environment is part of the worker's key, so two +actions that differ in it never share a process. `additional_environment` on +the config, the side-channel sources and the per-action `TMPDIR` are not +applied: they are per action, and the process is not. **A shared, long-lived working directory.** The process's own working directory is the worker's root action directory, not the per-action one; it has to be, @@ -6461,6 +6464,12 @@ ever serves. use it to validate or invalidate its own cached state. Anything the tool caches across requests, it caches on trust. +**Measured, not enforced.** What the process and its children use while they +serve a request is sampled and reported as the action's resource usage, so the +sizing loop learns what the tool needs. `resource_enforcement` never kills a +persistent worker for one request's reservation: the process is shared, and a +kill would take every action queued behind it. + The consequence is a real correctness requirement, not a footnote: **a persistent worker must be hermetic by construction, because nothing is enforcing it.** State that leaks between requests produces wrong outputs that are then @@ -6472,32 +6481,48 @@ it. ## The pool's actual behaviour [#the-pools-actual-behaviour] -Every parameter below is a compiled-in constant. None is configurable. +The pool is configured under `persistent_workers` on the worker; every value +below is its default. + +```json5 +persistent_workers: { + enabled: true, + max_workers_per_key: 4, + idle_timeout_s: 300, + max_requests_per_worker: 200, + shutdown_grace_ms: 5000, + acquire_timeout_s: 30, +} +``` -**Four workers per key.** A fifth concurrent action for the same key does not -queue; it falls back to normal one-shot execution and logs that it did. Your -build stays correct and quietly gets slower, which is worth knowing when you are -measuring. +**Four workers per key.** A fifth concurrent action for the same key waits +up to `acquire_timeout_s` for one of the four to come back, as Bazel's own +`worker_max_instances` queue does; only if none does in that time does it +fall back to normal one-shot execution, and it logs that it did. Your build +stays correct and gets slower by a tool startup per fallback; the worker's +`persistent_worker_fallbacks` counter says how often. **Two hundred requests per worker.** After that the process is retired and the next action starts a fresh one. This is deliberate recycling of accumulated JVM state, and it means a long build will restart workers periodically by design. -**No idle timeout, in practice.** The sweeper that would retire an idle worker -exists in the code but nothing calls it. A worker lives until the request cap -retires it, until it dies on its own, or until the NativeLink worker process -exits. +**Five minutes idle.** A sweeper runs on a timer and shuts down a worker +that has served nothing for `idle_timeout_s`, so a pool of JVMs left over from +one build does not hold the worker's memory through the next. **Dead workers are noticed on the next acquire**, not proactively. A JVM that OOMs while idle is discovered when the next action asks for it, and that action starts a replacement. -**Shutdown is not graceful.** Stdin is closed, giving the tool five seconds to -exit on EOF, and then it is killed outright. No `SIGTERM` is sent. +**Shutdown is not graceful.** Stdin is closed, giving the tool `shutdown_grace_ms` +to exit on EOF, and then it is killed outright. No `SIGTERM` is sent. **One request at a time per process.** Multiplexed workers are not supported; a request carrying a non-zero request ID is rejected. +**`enabled: false` turns it off.** Every action then runs one-shot, whatever it +asks for. + ## Steps [#steps] @@ -6585,10 +6610,13 @@ invalidating. -More than four concurrent actions for one key fall back to one-shot execution, -and one-shot execution of a worker-wrapped tool can be slower than the plain -tool. The per-key cap is not configurable, so the levers are reducing -concurrency for that mnemonic or splitting the work across more distinct keys. +More than `max_workers_per_key` concurrent actions for one key (four by +default) wait up to `acquire_timeout_s` for a process to come back and then +fall back to one-shot execution, and one-shot execution of a worker-wrapped +tool can be slower than the plain tool. Raise the cap, or reduce concurrency +for that mnemonic. The CPU time reported for a pooled action is what the +process spent on it; its peak memory is the whole process, including what +earlier requests left resident. @@ -8381,7 +8409,13 @@ client gives up. **`max_job_retries`** (default 3): how many times an action that fails with an internal error is retried before the last error is returned to the client. It exists to stop one poisonous action from cycling through and destabilising -your whole fleet. +your whole fleet. A lost worker and a memory escalation do not spend it. + +**`max_worker_loss_retries`** (default 10): how many times an action whose +worker was lost (disconnected, timed out, evicted, or OOM-killed as a whole) +is queued again before the client gets a `FAILED_PRECONDITION`. Losing a +worker is not the action's failure, so this budget is separate and larger; +it only stops an action that takes a worker down every time it runs. **`retain_completed_for_s`** (default 60 s): how long a completed action's result stays available for a late `WaitExecution` call. Clients that @@ -8392,6 +8426,119 @@ On the worker side, **`max_inflight_tasks`** (default 0, meaning unlimited) is the one to set on any real machine. Unlimited means the worker accepts every action the scheduler offers and lets the OS sort out the contention. +**`drain_on_shutdown`** (off by default) makes a worker announce that it is +draining the moment it gets `SIGTERM`, so the scheduler stops dispatching to +it while it finishes what it holds. It needs a scheduler that understands the +drain flag; an older one removes the worker on the spot and requeues its +actions. + +Three bounds keep one action from taking the worker with it. +**`max_captured_output_bytes`** (default 0, unbounded) caps how much of an +action's stdout or stderr the worker holds in memory; past it the output +spills to a file under the action directory and is uploaded from disk, so a +tool that prints gigabytes costs disk, not RAM. **`precondition_timeout_ms`** +(default 30000) ends a hanging precondition script and refuses the action as +backpressure. **`orphan_sweep_interval_s`** (default 0, off) removes action +directories under `work_directory` that no running action owns and that a +failed cleanup left behind; the directory is always purged at startup. + +When the worker ends an action, for a timeout or a scheduler cancel, the +process group gets SIGTERM first and **`kill_grace_ms`** (default 5000) to +write its own cleanup before SIGKILL follows; the memory reservation kill +skips the grace. A killed action keeps the stdout and stderr it produced. +**`set_tmpdir`** (default true) gives every action its own `TMPDIR` under its +action directory, so concurrent actions writing the same temporary file name +no longer collide; a `TMPDIR` in the action's own environment still wins, and +code that hard-codes `/tmp` is unaffected. A timeout above `max_action_timeout_s` +is clamped to it rather than rejected. + +When the action's own process has exited, anything it left running in its +process group (a background job, a helper that daemonized) is killed, since +it would otherwise hold the action's output pipes open and the action would +sit there until its timeout; the pipes then get five seconds to drain, and +what was read by then is kept. The worker also makes itself the parent of +every process an action leaves behind and reaps them on a timer, so a build +whose actions fork and do not wait no longer leaves hundreds of zombies on +the worker for the life of the pod. Under `use_namespaces` the stub does +the reaping inside the action's PID namespace. + +Every executed action reports its resource usage to the scheduler with an +`outcome`: `COMPLETED`, `KILLED_MEMORY` (the worker's reservation kill, or a +SIGKILL from the kernel with the last memory sample within 10% of the limit), +`KILLED_TIMEOUT` or `KILLED_EXTERNAL`, plus wall time and the reservation the +action ran under. The scheduler logs every kill outcome with its mnemonic. + +**`resource_enforcement`** (off by default, Linux only) makes the worker hold an +action to the memory reservation the scheduler placed it by. With +`memory: soft` the worker kills the action's process group once two +consecutive samples exceed `memory_kb` plus `memory_headroom_percent` +(default 20), or one sample exceeds twice that, and fails it with +`FAILED_PRECONDITION` naming the reservation and the observed peak. Samples +are 250 ms apart, so an action that allocates faster than the pod's whole +limit in half a second can still reach the cgroup limit first; the +reservations admitted onto a worker must also sum to less than its limit +with the headroom included. Without it the kill comes from the pod's cgroup limit, +and on cgroup v2 that kill takes the worker and every action on it, not just +the one that grew. The sample is the group's proportional set size, so a +parent's pages shared with its forked children count once. Actions run +through a persistent worker are not covered: that process outlives the action +and serves others, so only the pod's limit bounds it. With enforcement on, the +worker reads an action's properties from the scheduler's dispatch rather than +the client's request, so a reservation the scheduler placed the action by is +the one enforced and exported to the environment. + +`disk: guard` adds a check before any input is fetched: an action whose +`disk_kb` reservation exceeds the free space under the work directory, less +what the actions already admitted reserved, is refused with `RESOURCE_EXHAUSTED`, which the scheduler requeues without +counting an attempt and holds off this worker until its next keepalive, +instead of running the fetch into a full disk. +`disk: soft` is the guard plus a kill: every ten seconds the worker adds up +the files the action has written under its own directory (those modified +since its command started; inputs were materialized before that and are +not counted), and an action past its `disk_kb` reservation by +`disk_headroom_percent` is killed and fails with `FAILED_PRECONDITION`, its +outcome `KILLED_DISK`. The walk runs on the blocking pool as its own task, +so memory sampling carries on meanwhile. With either disk mode the action's +usage report carries `peak_disk_kb`, what it wrote, from one walk of the +action directory at completion, which is the only walk `guard` pays. The +figure is reported and exported for now; the scheduler's sizing consumer +comes with the size ladder. + +```json5 +"resource_enforcement": { + "memory": "soft", + "memory_property_name": "memory_kb", + "memory_headroom_percent": 20, + "disk": "soft", + "disk_property_name": "disk_kb", + "disk_headroom_percent": 20 +} +``` + +**`capacity`** makes the worker advertise CPU and memory from its own cgroup +limits instead of from `platform_properties`: it reads `cpu.max` and +`memory.max` at its cgroup v2 root, keeps back `overhead_cpu_millicores` +(thousandths of a core) and `overhead_memory_kb` for itself, divides the +memory by the enforcement headroom (`memory_headroom_percent`, which follows +`resource_enforcement` unless set), and sets the two properties to the +result at registration. The limit is the number the kernel enforces, so the +number the scheduler packs against is derived from it rather than typed next +to it. Where the cgroup cannot be read, hand-typed `cpu_count` and +`memory_kb` stand with a warning; with neither typed the worker refuses to +start, since it could take no action that asks for either. `cpu_unit` is the scale of the CPU +property: `cores` by default, rounded down, the scale `cpu_count: { +query_cmd: "nproc" }` and an action's `cpu_count=1` use; a fleet that +advertises and requests thousandths of a core sets `millicores`. The +overhead is always in thousandths of a core. + +```json5 +"capacity": { + "cpu_unit": "cores", + "overhead_cpu_millicores": 1000, + "overhead_memory_kb": 4194304 +} +``` + `experimental_backend` defaults to memory, so restarting the scheduler loses the queue. A Redis backend exists for shared state across scheduler @@ -13769,8 +13916,11 @@ workers: [{ `use_namespaces` puts each action in its own process namespace; `use_mount_namespace` additionally isolates the worker root in a new mount -namespace, and only works when `use_namespaces` is on. Both are off by default -and both are recommended in production. The container +namespace and gives each action a private `/tmp` (its own temporary directory, +bound over `/tmp` and removed with the action), and only works when +`use_namespaces` is on. Both are off by default and both are recommended in +production. Set `isolate_tmp: false` if actions must see the host's `/tmp`, +for example because a tool they run lives there. The container running the worker needs the privileges to create namespaces; the reference Compose stack marks the executor service `privileged: true` for exactly this. @@ -14255,8 +14405,9 @@ nativelink_executor: which need `CLONE_NEWUSER`/`CLONE_NEWNS`, hence the flag. Note what this buys and what it does not: namespaces here are a hermeticity and process-reaping mechanism, not a security boundary. Actions still run as the worker's own uid -through an identity uid map, still see the host filesystem, and still have -network access to every NativeLink port. A privileged container running +through an identity uid map, still see the host filesystem apart from a +private `/tmp`, and still have network access to every NativeLink port. A +privileged container running arbitrary build actions is a lateral-movement primitive. [Security hardening](/operate/security-hardening#the-sandbox-is-not-a-security-boundary) covers what the sandbox actually does. @@ -14578,10 +14729,14 @@ reconnect path calls `kill_all()` first, destroying every action it was running. Restart the scheduler when the queue is quiet, and restart it once. **Workers** drain themselves on `SIGTERM` (and on a first `SIGINT`; a second -`SIGINT` exits at once): the process rejects new actions, -waits for the ones it is running to finish, sends the scheduler a -`GoingAway`, and exits. `systemctl stop` is therefore safe as long as -`TimeoutStopSec` outlasts your longest action. To take a worker out of +`SIGINT` exits at once): with `drain_on_shutdown: true` +the process tells the scheduler at once that it is draining (a `GoingAway` +with the drain flag), so nothing new is dispatched to it, finishes the +actions it is running, and exits; the scheduler removes it when the stream +closes. Without the flag the worker finishes first and announces itself +last, and until then the scheduler keeps offering it work it refuses. The +wait is bounded by `max_action_timeout_s`. `systemctl stop` is therefore +safe as long as `TimeoutStopSec` outlasts your longest action. To take a worker out of rotation without stopping it, and if your scheduler configuration enables the `admin` service on the worker API listener: @@ -14949,7 +15104,8 @@ each one decides where you have to point the collector: | Cache | `nativelink_cache_*` | Any process with a store wrapped in `cache_metrics` | Hit rate, operation latency, bytes in and out, per cache type | | Execution | `nativelink_execution_*` | The scheduler | Queue time, stage durations, active and completed counts, retries, and per-action CPU time and peak memory as reported by the worker that ran it | | Worker fleet | `nativelink_worker_*` | The scheduler | Workers connected, joins and departures by reason, keepalives, and how many are paused or draining | -| Scheduler | `nativelink_scheduler_matching_*` | The scheduler | How long each matching pass takes and whether it succeeded | +| Scheduler | `nativelink_scheduler_matching_*` | The scheduler | How long each matching pass takes, whether it succeeded, and how many waiting actions took room that opened mid-pass ahead of newer ones | +| Sizing | `nativelink_scheduler_hints_resolutions_total` | The scheduler, with `historical_resource` configured | Actions sized, by source: a hint (`target`, `action_digest`, `command_digest`, `mnemonic`), `cold_start`, or `none`; the `cold_start` share is what a hint feed is meant to drive down | | gRPC server | `nativelink_rpc_server_duration_*` | Every process, for every gRPC service it serves | Request latency by service, method and status code | | Store tier | `nativelink_store_tier_*` | Any process with a `fast_slow` store | Fast and slow tier hits, misses and stale reads, and bytes read from each tier | | Health | `nativelink_health_checks_*` | Every process, on each run of the health endpoint | Health check outcomes per component: `ok`, `initializing`, `warning`, `failed` or `timeout` | @@ -15002,6 +15158,16 @@ sum(rate(nativelink_execution_completed_count_total{execution_result="success"}[ sum(rate(nativelink_execution_completed_count_total[5m])) ``` +Actions the worker killed for memory, as a share of the sampled ones. The +per-action resource histograms carry `execution_outcome` (`completed`, +`killed_memory`, `killed_disk`), so a kill can be told from a completion and +the memory escalation rate read against it: + +```text +sum(rate(nativelink_execution_peak_memory_count{execution_outcome="killed_memory"}[5m])) / +sum(rate(nativelink_execution_peak_memory_count[5m])) +``` + Workers connected to each scheduler instance, and how many of them are paused or draining: @@ -15129,10 +15295,10 @@ not your pipeline. Queue time, not queue depth. `nativelink_execution_queue_time` is a real histogram emitted by the scheduler; a p95 above your tolerance means add - workers. No `nativelink_scheduler_queue_depth` series exists; the - closest thing is `nativelink_execution_active_count{execution_stage="queued"}`, - a gauge, and `nativelink_worker_connected_count` tells you how many workers - that queue is waiting on. Before you build an HPA on either, read the + workers. `nativelink_scheduler_queue_depth` is a gauge of what the last + matching pass read, `nativelink_execution_active_count{execution_stage="queued"}` + counts the same thing from the other side, and `nativelink_worker_connected_count` + tells you how many workers that queue is waiting on. Before you build an HPA on either, read the [autoscaling reference](/operate/autoscaling-reference#the-reaction-time-budget): the default one-minute export interval alone puts your reaction time in minutes, and CPU-based scaling is worse because workers are mostly I/O-bound. @@ -15481,9 +15647,14 @@ uid-remapping sandbox; it is a process-reaping and hermeticity mechanism. PID namespacing means orphaned children die with the action instead of leaking; IPC and UTS namespacing means actions cannot see each other's semaphores or hostname (the hostname is set to `nativelink`). With `use_mount_namespace` -the flag list also gains `CLONE_NEWNS`, and the worker's own root action -directory is masked so one action cannot see another's tree; the rest of the -host filesystem stays visible. +the flag list also gains `CLONE_NEWNS`, the worker's own root action +directory is masked so one action cannot see another's tree, and `/tmp` is +replaced with the action's own temporary directory, on the worker's disk and +removed with the action, so actions cannot collide on, or leak state +through, predictable paths there (`isolate_tmp: false` keeps the host's +`/tmp` visible; the bound `/tmp` carries the work volume's mount options, +`noexec` included, not the host `/tmp`'s); the rest of the host filesystem +stays visible. Note what is absent from that flag list: **`CLONE_NEWNET`**. No network namespace is created. An action has the worker's full network access, including to port @@ -15813,6 +15984,22 @@ exist at all. Nothing in the repository wires this up for you. It is a short script, and it's the difference between a clean scale-in and a burst of retried actions. +Two read-only listings sit next to it, for the same scripts and for an +autoscaler that wants the scheduler's own view rather than its store: + +``` +GET {admin_path}/scheduler/{scheduler_name}/workers +GET {admin_path}/scheduler/{scheduler_name}/demand +``` + +`workers` returns every connected worker with what it registered, what it +has left after its running actions' reservations, its running count, its +pause and drain flags and the free memory it last reported. `demand` +returns every queued action with its operation id, when it was queued, its +priority and its platform properties, reservations included. A worker's id +starts with the name it registered under, which on Kubernetes is the pod +name, so a drain script can find the ids to drain from the pod. + @@ -16698,6 +16885,43 @@ each check) and `200` otherwise. "nothing objected," which for most compositions is vacuous. +## Readiness is a different question [#readiness-is-a-different-question] + +`/status` answers whether the components that registered a health indicator +have failed. It stays `200` while a component is still initializing, which is +the right answer for a liveness probe and the wrong one for readiness. The +`health` service therefore also serves `/ready` (`readiness_path`, default +`/ready`), which answers `503` while anything is initializing or failed and +`200` otherwise. + +For a worker, the component that matters is its registration with the +scheduler. Each worker registers a `workers/` indicator that reads +`Initializing` until the scheduler has answered its `connect_worker` with a +`ConnectionResult`, `Ok` while that connection holds, and `Initializing` +again after it is lost and the worker is reconnecting, so only `/ready` +drops: a liveness probe on `/status` stays green through a scheduler roll +instead of restarting every worker. A Kubernetes readiness probe on the +worker's `/ready` turns the pod Ready when the worker can take work rather +than when its container started, and turns it not Ready again if it loses +the scheduler. Anything that reads pod readiness, an autoscaler's idle clock +included, then counts from registration. + +```json5 +// On the worker pod's HTTP listener. +services: { + health: { path: "/status", readiness_path: "/ready" }, +} +``` + +```yaml +readinessProbe: + httpGet: { path: /ready, port: 50061 } + periodSeconds: 5 + failureThreshold: 3 +``` + + + ## Build the deployment [#build-the-deployment] @@ -17290,15 +17514,16 @@ signal deaths, so the log line is what you alert on. **The worker process itself was killed.** The scheduler notices after `worker_timeout_s` (default 10 seconds, checked once a second), evicts it, and -re-queues its in-flight actions. In this tree a disconnect *counts as an -attempt*, so after `max_job_retries` the client finally sees, as a +re-queues its in-flight actions. A lost worker does not spend the action's +`max_job_retries`; it has its own budget, `max_worker_loss_retries` (default +10), so after that many losses the client finally sees, as a `FAILED_PRECONDITION` that Bazel does not resubmit: ```text -Worker disconnected repeatedly while executing this action (4 > 3 attempts); -the worker was likely OOM-killed or its pod evicted. Give the action a memory -reservation or raise the pool's memory limit before retrying. for operation_id: ..., -maybe_worker_id: ... +The worker running this action was lost 11 times, more than +max_worker_loss_retries (10); the action is likely what takes the worker down +(OOM-killed pod, eviction). Give it a memory reservation or raise the pool's +memory limit before retrying. for operation_id: ..., maybe_worker_id: ... ``` That string is the single best diagnostic for a memory-pressured worker fleet. @@ -17587,6 +17812,85 @@ The new open file limit (1024) is below the recommended value of 24576. Consider raising max_open_files. ``` +## Every build stalls at `[Sched]` and nothing errors [#every-build-stalls-and-nothing-errors] + +Clients sit on `WaitExecution` for hours, the scheduler log is quiet apart +from `Could not find AwaitedAction for the given operation id`, idle workers +take nothing, and the provisioner reads demand as zero. The scheduler's +records are being evicted from Redis. + +`maxmemory` and `maxmemory-policy` are per Redis server, not per database. +A deployment that puts the scheduler's records in the same instance as the +CAS's small blobs and the BEP event stream, under `allkeys-lru`, has agreed +that growth in any of them may evict any of them. An evicted CAS blob is a +cache miss. An evicted scheduler record is a client waiting on an operation +nothing will ever finish, and a queue entry that, until v1.7.3, stopped the +abandoned sweep at the first record it could not read, so the corpses behind +it were never retired either. + +What to check, in order: + +1. `redis-cli info memory`: `used_memory` against `maxmemory`, and + `evicted_keys` in `info stats`. Any eviction at all on the scheduler's + instance is this incident in progress. +2. What is growing. A stalled BEP consumer is the usual cause: count keys by + prefix. Postgres disk full stops the consumer, so check that too. +3. The scheduler's own counters: `scheduler.awaited_action.orphans` (queue + entries whose record was gone on read), `scheduler.queue.depth`, + `scheduler.sweep.failures`, and `worker.connected.count`. Depth above + zero with no connected worker for minutes is the stall itself. + +What to change: + +- Give the scheduler its own Redis with `maxmemory-policy noeviction`. A + full instance then refuses writes, which is an error the scheduler reports + and a client sees, instead of a silent loss. The Helm chart deploys one as + `scheduler-redis`. +- Bound the BEP stream with `key_ttl_s` sized for the longest consumer + outage you mean to survive, not for convenience. +- Alert on Redis memory and evictions and on the counters above before the + queue does the alerting for you. + +From v1.7.3 the matching pass and the abandoned sweep skip a listed +operation whose record is gone, count it, and carry on, so an eviction no +longer makes the pile permanent. It still makes every affected client wait +out its timeout. + +## The scheduler's Redis keeps every completed action [#scheduler-redis-keeps-completed-actions] + +**Symptom.** With the Redis scheduler backend, `DBSIZE` on the scheduler's +database grows by one `aa_*` hash per action ever run and never falls; the +records all read `completed` and `TTL` says `-1`. Once the search index holds +more than one cursor page, listings over it (the admin read the worker +provisioner makes, the scheduler admin endpoints) can fail with +`Failed to get operations list! Parse: Non map item`. + +**Cause.** Before v1.7.3 the Redis backend received `retain_completed_for_s` +as configured, without the default applied: left unset it arrived as 0, and the versioned update treats a +zero TTL as no expiry, so completed records were written to live forever. +The memory backend applied the 60 second default; only Redis leaked. The +listing failure is a second bug in the same release: a cursor page arrived +in the RESP3 shape the first page's parser understood but the cursor reader +did not. + +**Fix.** Upgrade to v1.7.3 or later; new completions expire after +`retain_completed_for_s`. Records written before the upgrade keep their +missing expiry, so give them one once, on the scheduler's Redis database +(0 in the chart's layout): + +```bash +redis-cli -n 0 --scan --pattern 'aa_*' \ + | while read -r key; do + if [ "$(redis-cli -n 0 hget "$key" state)" = "completed" ] \ + && [ "$(redis-cli -n 0 ttl "$key")" = "-1" ]; then + redis-cli -n 0 expire "$key" 60 > /dev/null + fi + done +``` + +A queued or executing record is left alone; only completed ones get the +sixty seconds a fresh completion would. + ## Restarting the scheduler [#restarting-the-scheduler] Not an incident so much as a procedure that behaves worse than you expect. @@ -18964,11 +19268,78 @@ values per worker. For clusters of tens or low hundreds of workers this is not the bottleneck. It is worth knowing about before you plan for thousands. -`allocation_strategy` is a smaller knob than its name suggests: it chooses -**which end of the LRU the scan starts from**, and nothing else. Starting -from the least-recently-used end spreads load; starting from the -most-recently-used end packs work onto hot workers and leaves the rest idle -long enough to be scaled down. Neither is a bin-packing algorithm. +`allocation_strategy` is a smaller knob than its name suggests. The two +original values choose **which end of the LRU the scan starts from**, and +nothing else. Starting from the least-recently-used end spreads load; +starting from the most-recently-used end packs work onto hot workers and +leaves the rest idle long enough to be scaled down. `least_loaded` takes the +eligible worker with the fewest running actions, least recently used among +equals, so a worker that just joined takes the next action instead of +waiting for every older worker to be used once more, and a burst that lands +while most of the fleet is still starting does not fill the first worker to +its concurrency cap. None of them is a bin-packing algorithm. + +The admission ledger only subtracts what an action declares, so an action +that declared too little is invisible to it. **`live_memory_veto`** names a +`minimum` property, usually `memory_kb`, that the scheduler checks against the +free memory each worker reports with its keepalive (the cgroup limit less +current usage, or `MemAvailable` without a limit). A worker reporting less +than the action asks for is skipped for that action; a worker that reports +nothing, such as one older than v1.7.3, is never vetoed. Two metrics go with +this: `worker.keepalive.gap` is the seconds between consecutive messages from +a worker, which is what `worker_timeout_s` should be judged against, and +`worker.disconnections` carries a `worker.disconnect.reason` of +`disconnected`, `drained`, `timeout`, `kill_unacknowledged`, `error`, +`removed` or `shutdown`. + +Within a priority the queue is ordered by insert time to the nanosecond, on +both backends, so a burst submitted together dispatches in arrival order. + +The pass keeps that order when a worker frees up while it is running. A pass +reads the queue one action at a time, a store round trip each on a +store-backed queue (the listing serves the record it already loaded, and +only the client keepalive is read), so on a long queue a worker that +finishes mid-pass would otherwise hand its room to whatever action the pass +happened to be reading, newer than everything it had already passed over. Instead, actions that found +no room are parked in the order the queue listed them, highest priority first +and then oldest, and when the fleet's capacity changes the pass offers the +room to the parked actions in that order before the one in hand, asking the +fleet once per property shape. `scheduler.matching.parked_dispatched` counts +the parked actions a worker was actually told to run that way; a non-zero +rate is normal on a busy fleet and means the head of the queue is being +served. + +When a worker reports an action killed for memory, by its own reservation +enforcement or by the kernel near the limit, **`memory_escalation`** requeues +the action with the reservation multiplied by `percent` (default 200) instead +of failing it. With `ladder_kb` set to the memory values of the fleet's size +classes, a kill steps the reservation to the next class up instead of +scaling it. The next dispatch carries the raised value, so the worker holds +the action to that. The reservation grows to the largest memory any connected +worker advertises (or `max_kb`); the last step reserves that worker whole, +its memory and, through `cpu_property`, its CPU, so the action runs alone. +Only a kill there fails the action, with a message saying no worker in the +fleet can run it. Escalations have their own budget (`max_steps`, default 8) +and do not spend `max_job_retries`. The execution histograms carry an +`execution.outcome` attribute so the killed samples can be told from the +completed ones. + +The `historical_resource` scheduler that fills reservations reads hints by a +key ladder, each step probed only when the action carries that key: +`(target_id, action_mnemonic)`, then `target_id`, then the +action's `action_digest`, then its `command_digest` (each `hash-size` or +`hash/size`), then the mnemonic alone. The digest keys reach every client, +whether or not it sends Bazel request metadata; the action digest is what +origin events carry, so it is what a producer fed by them writes. A hint may name a `class` instead of numbers, and +**`classes`** on the same scheduler defines the fleet's size ladder: named +points with `cpu_count`, `memory_kb` and `disk_kb`, listed ascending. An +action no hint matches takes the class picked by `class_by_mnemonic`, then +`class_by_timeout` (Bazel test sizes arrive as timeouts), then +`default_class`; without classes the raw `cold_start` numbers apply. With +`memory_escalation.ladder_kb` set to the classes' memory values, a memory +kill steps the reservation to the next class instead of scaling it, and a +kill at the top class fails the action. Older hint files still load: every +new field is optional. ### Platform properties are typed, and `Minimum` is a resource pool [#platform-properties-are-typed-and-minimum-is-a-resource-pool] @@ -19005,17 +19376,25 @@ the configuration side. ## Retries, and what "failed" means [#retries-and-what-failed-means] -The attempt counter is incremented by two different events, with different -rules. +The attempt counter is incremented by one event, and two other events that +used to spend it now have budgets of their own. **`UpdateWithError` increments the attempt count, unless the error is `ResourceExhausted`.** Backpressure is free. A worker saying "I am too busy" does not spend one of the action's lives, which is what makes it safe for a worker to shed load. -**`UpdateWithDisconnect` always increments.** A worker that vanishes mid-action -always costs an attempt. Without this, an action that reliably OOM-kills its -worker would retry forever, taking a worker down each time. +**`UpdateWithDisconnect` counts against `max_worker_loss_retries`** (default +10), not `max_job_retries`. A worker that vanishes mid-action is usually not +the action's doing, so the action goes back to the queue without spending one +of its lives; the separate budget is what stops an action that reliably +OOM-kills its worker from taking one down forever. Past it the client sees a +`FAILED_PRECONDITION` saying the worker was lost that many times. + +**`UpdateWithEscalation` counts against `memory_escalation.max_steps`** +(default 8), not `max_job_retries` either: the scheduler chose to run the +action again with more memory, and the fleet's ceiling bounds how far that +goes. The limit is checked as `attempts > max_job_retries`, strictly, so the default of `3` permits **four** total executions. That off-by-one is in the @@ -19293,7 +19672,8 @@ command finished slightly before the bytes exist. participant sched as Scheduler participant w as Worker participant cas as CAS + AC - sched->>w: assign action (over ConnectWorker) + sched->>w: StartExecute (over ConnectWorker) + w->>sched: ExecuteAccepted, or ExecuteDeclined and stop here w->>w: create work directory for the operation w->>cas: fetch missing blobs w->>w: hardlink inputs from filesystem CAS @@ -19305,6 +19685,24 @@ command finished slightly before the bytes exist. w->>sched: final result update `} +The first reply is the worker's admission decision. A dispatch is a +proposal until the worker answers it: `ExecuteAccepted` confirms the ledger +charge the scheduler made on the send, `ExecuteDeclined` hands it back. The +worker declines when it already holds `max_inflight_tasks` actions, when the +action's memory reservation is more than the worker currently has free, or +when it is shutting down. A declined action returns to the queue without +counting as an attempt, and the worker is paused until its next keepalive, +or, for a decline for memory, until a keepalive reports that much free. The +scheduler's channel to each worker is bounded at its concurrency plus a +margin; a worker that has stopped reading has its dispatch requeued the same +way rather than piling up work it will never see. A worker only sends these +replies to a scheduler that announced support in its `ConnectionResult`, so +an older scheduler sees the old behaviour: the worker runs whatever it is +sent. `dispatch_ack_timeout_s` on the scheduler adds the converse: a dispatch +no worker acknowledged within the timeout is taken back, which only makes +sense once every connected worker acknowledges. + + Four behaviours in the upload path: **A declared-but-missing output is not an error.** An action that promises @@ -19499,8 +19897,11 @@ is the toolchain half of the answer. ### Actions can run more than once [#actions-can-run-more-than-once] The default `max_job_retries` of 3 is checked strictly, so an action may -execute **four** times. A worker disconnect always costs an attempt; a -`ResourceExhausted` never does. +execute **four** times on its own failures. A lost worker has its own budget, +`max_worker_loss_retries` (default 10), and a memory escalation its own +(`memory_escalation.max_steps`, default 8), so an action can run more often +than four times when workers die under it or it keeps outgrowing its +reservation; a `ResourceExhausted` never costs anything. For a pure function this is invisible. For an action with side effects (one that publishes, deploys, increments, or writes anywhere outside its declared @@ -23132,8 +23533,10 @@ a how-to showing a renamed field keeps showing it, and nothing fails until a reader copies it and the binary answers `unknown field`. `lint:snippets` closes that gap. It harvests every field name out of the -generated reference and checks every `json`/`json5` fence under -`content/docs/` against it. +generated reference, the latest release's page and main's, and checks every +`json`/`json5` fence under `content/docs/` against it. A snippet for a field +added since the last release passes once `gen:config-reference main` has +been run in the same change. ```bash bun --filter @nativelink/docs lint:snippets diff --git a/llms-small.txt b/llms-small.txt index 2510d675a48..a362db40db4 100644 --- a/llms-small.txt +++ b/llms-small.txt @@ -1653,11 +1653,12 @@ reusing a warm process, and a clear-eyed view of the isolation you trade for it. ## Two things to know before you start -**Almost nothing needs configuring in NativeLink.** No field anywhere in -[the configuration model](/configuration) turns this on, tunes the pool, or -selects a protocol; `LocalWorkerConfig` has no persistent-worker section at -all, and every pool parameter is a compiled-in constant. The feature is -driven by two platform properties on each action. +**Little needs configuring in NativeLink.** The pool is on by default and +the protocol is chosen by two platform properties on each action. The +worker's `persistent_workers` block (v1.7.3 and later) sizes it: how many +processes per key, how long an idle one lives, how many requests one serves, +how long an action at the cap waits for a process to come back, and +`enabled: false` to run every such action one-shot. ## Opt an action in @@ -1672,13 +1673,14 @@ Actions are pooled by a key derived from the command line: ## The isolation you give up -This is the part to read twice. NativeLink applies real isolation to ordinary -actions and **applies none of it to persistent workers**. The persistent path -returns before any of it runs. +This is the part to read twice. A persistent worker process outlives the +actions it serves, so some of the isolation an ordinary action gets cannot +apply to it, and some of it did not apply until v1.7.3. ## The pool's actual behaviour -Every parameter below is a compiled-in constant. None is configurable. +The pool is configured under `persistent_workers` on the worker; every value +below is its default. ## Steps @@ -3457,6 +3459,15 @@ short fix and all five are prerequisites. Before you write a readiness probe, understand what `/status` actually measures, because the answer is surprising in both directions. +## Readiness is a different question + +`/status` answers whether the components that registered a health indicator +have failed. It stays `200` while a component is still initializing, which is +the right answer for a liveness probe and the wrong one for readiness. The +`health` service therefore also serves `/ready` (`readiness_path`, default +`/ready`), which answers `503` while anything is initializing or failed and +`200` otherwise. + ## Build the deployment
  • @@ -3596,6 +3607,22 @@ in response to what the worker returns. Several NativeLink limits queue instead of erroring, which is why "hung" is a more common symptom than "failed": +## Every build stalls at `[Sched]` and nothing errors + +Clients sit on `WaitExecution` for hours, the scheduler log is quiet apart +from `Could not find AwaitedAction for the given operation id`, idle workers +take nothing, and the provisioner reads demand as zero. The scheduler's +records are being evicted from Redis. + +## The scheduler's Redis keeps every completed action + +**Symptom.** With the Redis scheduler backend, `DBSIZE` on the scheduler's +database grows by one `aa_*` hash per action ever run and never falls; the +records all read `completed` and `TTL` says `-1`. Once the search index holds +more than one cursor page, listings over it (the admin read the worker +provisioner makes, the scheduler admin endpoints) can fail with +`Failed to get operations list! Parse: Non map item`. + ## Restarting the scheduler Not an incident so much as a procedure that behaves worse than you expect. @@ -3958,8 +3985,8 @@ Every property key a worker advertises has a declared type: ## Retries, and what "failed" means -The attempt counter is incremented by two different events, with different -rules. +The attempt counter is incremented by one event, and two other events that +used to spend it now have budgets of their own. ### Two timeouts with very different defaults @@ -4100,8 +4127,11 @@ result computed with a different value of it. ### Actions can run more than once The default `max_job_retries` of 3 is checked strictly, so an action may -execute **four** times. A worker disconnect always costs an attempt; a -`ResourceExhausted` never does. +execute **four** times on its own failures. A lost worker has its own budget, +`max_worker_loss_retries` (default 10), and a memory escalation its own +(`memory_escalation.max_steps`, default 8), so an action can run more often +than four times when workers die under it or it keeps outgrowing its +reservation; a `ResourceExhausted` never costs anyth… ### Some output metadata is discarded diff --git a/nativelink-config/src/cas_server.rs b/nativelink-config/src/cas_server.rs index 87f186d57d2..3e52a8d1741 100644 --- a/nativelink-config/src/cas_server.rs +++ b/nativelink-config/src/cas_server.rs @@ -449,6 +449,17 @@ pub struct HealthConfig { /// Timeout on health checks. Default: 5s. #[serde(default)] pub timeout_seconds: u64, + + /// Path of the readiness check, the stricter sibling of `path`: it + /// answers 503 while any component is still initializing, where `path` + /// answers 200. A worker's registration with its scheduler is such a + /// component, so a Kubernetes readiness probe on this path turns Ready + /// only once the worker can take work. Must differ from `path`; the + /// same path for both is refused at startup. + /// + /// Default: "/ready" + #[serde(default)] + pub readiness_path: String, } #[derive(Deserialize, Serialize, Debug)] @@ -743,7 +754,106 @@ pub struct ServerConfig { pub experimental_identity_header: IdentityHeaderSpec, } -#[derive(Deserialize, Serialize, Debug)] +/// The scale a CPU property is advertised on. +#[derive(Copy, Clone, Deserialize, Serialize, Debug, Default, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +#[cfg_attr(feature = "dev-schema", derive(JsonSchema))] +pub enum CpuUnit { + /// Whole cores, rounded down: what `cpu_count: { query_cmd: "nproc" }` + /// advertises and what an action asking for `cpu_count=1` means. + #[default] + Cores, + /// Thousandths of a core, for a fleet whose actions and workers already + /// speak that scale. + Millicores, +} + +/// Advertise CPU and memory from what the worker can actually see. The +/// worker reads `cpu.max` and `memory.max` from its own cgroup v2 directory +/// (found through `/proc/self/cgroup`, so a privileged container that sees +/// the host's tree still finds its own) up to the nearest limited ancestor, +/// the pod's limits on Kubernetes; with no limit at any level the configured +/// properties stand. It takes off what it needs for itself, divides +/// the memory by the enforcement headroom, and sets the two properties to +/// the result at registration. What the scheduler packs against is then +/// derived from the one number that is enforced, instead of typed in twice. +#[derive(Deserialize, Serialize, Debug, Clone)] +#[serde(deny_unknown_fields)] +#[cfg_attr(feature = "dev-schema", derive(JsonSchema))] +pub struct CapacityConfig { + /// Property that carries CPU, on the scale `cpu_unit` names. + /// Default: `cpu_count` + #[serde( + default = "default_capacity_cpu_property_name", + deserialize_with = "convert_string_with_shellexpand" + )] + pub cpu_property_name: String, + + /// Property that carries memory, in KiB. + /// Default: `memory_kb` + #[serde( + default = "default_capacity_memory_property_name", + deserialize_with = "convert_string_with_shellexpand" + )] + pub memory_property_name: String, + + /// The scale the CPU property is advertised on. Whole cores by default, + /// the scale every other example in this configuration uses; a fleet + /// that advertises and requests thousandths of a core sets `millicores`. + /// Getting this wrong is a thousandfold error in how many actions the + /// scheduler packs onto the worker. + /// Default: cores + #[serde(default)] + pub cpu_unit: CpuUnit, + + /// CPU the worker keeps for itself, in thousandths of a core whatever + /// `cpu_unit` says. + /// Default: 1000 + #[serde( + default = "default_capacity_overhead_cpu_millicores", + deserialize_with = "convert_numeric_with_shellexpand" + )] + pub overhead_cpu_millicores: u64, + + /// Memory the worker keeps for itself (the process, its directory cache, + /// output buffers), in KiB. + /// Default: 2097152 (2 GiB) + #[serde( + default = "default_capacity_overhead_memory_kb", + deserialize_with = "convert_numeric_with_shellexpand" + )] + pub overhead_memory_kb: u64, + + /// Percent the advertised memory is reduced by, so that actions at their + /// enforcement ceiling (reservation plus headroom) still fit the limit. + /// Unset, it follows `resource_enforcement.memory_headroom_percent` + /// while memory enforcement is on, and is 0 otherwise, so the two + /// cannot drift apart by being typed twice. + /// Default: unset + #[serde( + default, + deserialize_with = "convert_optional_numeric_with_shellexpand" + )] + pub memory_headroom_percent: Option, +} + +fn default_capacity_cpu_property_name() -> String { + "cpu_count".to_string() +} + +fn default_capacity_memory_property_name() -> String { + "memory_kb".to_string() +} + +const fn default_capacity_overhead_cpu_millicores() -> u64 { + 1000 +} + +const fn default_capacity_overhead_memory_kb() -> u64 { + 2_097_152 +} + +#[derive(Deserialize, Serialize, Debug, Clone)] #[serde(rename_all = "snake_case")] #[cfg_attr(feature = "dev-schema", derive(JsonSchema))] pub enum WorkerProperty { @@ -775,6 +885,108 @@ pub struct EndpointConfig { pub tls_config: Option, } +#[derive(Copy, Clone, Deserialize, Serialize, Debug, Default, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +#[cfg_attr(feature = "dev-schema", derive(JsonSchema))] +pub enum MemoryEnforcement { + /// Measure only. + None, + /// Kill the action's process group once two consecutive samples exceed + /// the reservation plus headroom, or one sample exceeds twice it. + #[default] + Soft, +} + +#[derive(Clone, Deserialize, Serialize, Debug)] +#[serde(deny_unknown_fields)] +#[cfg_attr(feature = "dev-schema", derive(JsonSchema))] +pub struct ResourceEnforcementConfig { + /// What to do with an action that exceeds its memory reservation. + /// Default: soft + #[serde(default)] + pub memory: MemoryEnforcement, + + /// Whether to refuse an action whose disk reservation exceeds the free + /// space under the work directory, less what the actions already + /// admitted reserved, before its inputs are fetched. The + /// refusal is `ResourceExhausted`, which the scheduler requeues without + /// counting an attempt and holds off this worker until its next + /// keepalive. + /// Default: none + #[serde(default)] + pub disk: DiskEnforcement, + + /// Platform property carrying the action's disk reservation in KiB. + /// Default: `disk_kb` + #[serde( + default = "default_disk_property_name", + deserialize_with = "convert_string_with_shellexpand" + )] + pub disk_property_name: String, + + /// Platform property carrying the action's memory reservation in KiB, + /// as the scheduler sends it. An action without the property is not + /// enforced. + /// Default: `memory_kb` + #[serde( + default = "default_memory_property_name", + deserialize_with = "convert_string_with_shellexpand" + )] + pub memory_property_name: String, + + /// Percent above the reservation an action may reach before it is + /// killed. + /// Default: 20 + #[serde( + default = "default_memory_headroom_percent", + deserialize_with = "convert_numeric_with_shellexpand" + )] + pub memory_headroom_percent: u64, + + /// Percent above its disk reservation an action's own files may reach + /// before it is killed, with `disk: soft`. + /// Default: 20 + #[serde( + default = "default_memory_headroom_percent", + deserialize_with = "convert_numeric_with_shellexpand" + )] + pub disk_headroom_percent: u64, +} + +fn default_memory_property_name() -> String { + "memory_kb".to_string() +} + +const fn default_true() -> bool { + true +} + +fn default_disk_property_name() -> String { + "disk_kb".to_string() +} + +#[derive(Copy, Clone, Deserialize, Serialize, Debug, Default, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +#[cfg_attr(feature = "dev-schema", derive(JsonSchema))] +pub enum DiskEnforcement { + /// Fetch inputs regardless of free space. + #[default] + None, + /// Refuse the action when the free space under the work directory is + /// below its reservation. + Guard, + /// `guard`, and kill an action whose files under its own directory + /// grow past its reservation plus `disk_headroom_percent`. Sampled + /// every ten seconds by walking the action directory, counting only + /// files the action itself wrote (inputs are hard links into the CAS + /// and are not counted). + Soft, +} + +const fn default_memory_headroom_percent() -> u64 { + 20 +} + #[derive(Copy, Clone, Deserialize, Serialize, Debug, Default)] #[serde(rename_all = "snake_case")] #[cfg_attr(feature = "dev-schema", derive(JsonSchema))] @@ -905,6 +1117,68 @@ pub struct UploadActionResultConfig { pub failure_message_template: String, } +/// The pool of persistent worker processes (Bazel `supports-workers`). +#[derive(Deserialize, Serialize, Debug, Clone, Copy)] +#[serde(deny_unknown_fields)] +#[cfg_attr(feature = "dev-schema", derive(JsonSchema))] +pub struct PersistentWorkersConfig { + /// Run an action whose platform properties carry `supports-workers=1` + /// in a pooled worker process that outlives it and serves the next + /// action with the same executable, startup arguments and environment. + /// Off, every such action runs as a one-shot process like any other. + /// Default: true + #[serde(default = "default_true")] + pub enabled: bool, + + /// Most worker processes kept per key (executable, startup arguments, + /// environment, protocol). An action whose key is at the cap and has + /// no idle process waits `acquire_timeout_s` for one to come back, then + /// runs one-shot. 0 takes the default. + /// Default: 4 + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub max_workers_per_key: usize, + + /// Seconds an idle worker process is kept before the sweeper shuts it + /// down. 0 takes the default. + /// Default: 300 + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub idle_timeout_s: u64, + + /// Requests a worker process serves before it is retired and replaced, + /// which bounds what a long-lived process accumulates. 0 takes the + /// default. + /// Default: 200 + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub max_requests_per_worker: u64, + + /// Milliseconds a worker process being shut down gets to exit on its + /// own (its stdin is closed) before SIGKILL. 0 takes the default. + /// Default: 5000 + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub shutdown_grace_ms: u64, + + /// Seconds an action waits for a worker process to come back when its + /// key is at `max_workers_per_key` with none idle, before it runs + /// one-shot instead. 0 takes the default. + /// Default: 30 + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub acquire_timeout_s: u64, +} + +/// The same as an absent block: the pool on, every size at its default. +impl Default for PersistentWorkersConfig { + fn default() -> Self { + Self { + enabled: true, + max_workers_per_key: 0, + idle_timeout_s: 0, + max_requests_per_worker: 0, + shutdown_grace_ms: 0, + acquire_timeout_s: 0, + } + } +} + /// Opt-in file capture for the actual container executing a Buck2 action. /// Requires a fresh single-use container with no outer execution wrapper or /// separate action mount namespace. For nested runtimes, run `NativeLink` and the @@ -952,6 +1226,10 @@ pub struct Buck2FileCaptureConfig { pub finalize_timeout_s: usize, } +#[expect( + clippy::struct_excessive_bools, + reason = "independent on/off switches of one worker; a nested struct per flag would not read better" +)] #[derive(Deserialize, Serialize, Debug, Default)] #[serde(deny_unknown_fields)] #[cfg_attr(feature = "dev-schema", derive(JsonSchema))] @@ -989,6 +1267,21 @@ pub struct LocalWorkerConfig { )] pub max_upload_timeout_s: usize, + /// Maximum time allowed for fetching an action's inputs into the worker's + /// store and its directory before the command runs. A fetch that hangs + /// (a store that never answers) otherwise holds the action's slot for + /// good while the worker looks healthy; past this the action fails with + /// `DeadlineExceeded` and may be retried by the scheduler, and a warning + /// names the action every minute before that. Value in seconds. + /// + /// Default: 10 minutes + #[serde( + default, + deserialize_with = "convert_duration_with_shellexpand", + alias = "max_download_timeout" + )] + pub max_download_timeout_s: usize, + /// Maximum time to wait for action directory cleanup before timing out. /// Value in seconds. /// @@ -1047,6 +1340,31 @@ pub struct LocalWorkerConfig { #[serde(default)] pub timeout_handled_externally: bool, + /// On shutdown, tell the scheduler at once that this worker is draining + /// so it stops dispatching here, then finish the running actions (up to + /// `max_action_timeout_s`) and exit. Off, the worker finishes first and + /// only then announces itself, while the scheduler keeps offering it + /// work it has to refuse. Needs a scheduler that understands the drain + /// flag (v1.7.3+); an older one removes the worker and requeues its + /// actions as soon as the message arrives. + /// Default: false + #[serde(default)] + pub drain_on_shutdown: bool, + + /// Kill an action that grows past its memory reservation before the + /// pod's cgroup limit does. On cgroup v2 the container cgroup carries + /// `memory.oom.group`, so the kernel's kill takes the worker and every + /// action on it; this one takes only the offender and fails it with + /// `FailedPrecondition`. Linux only; accepted and ignored elsewhere. + /// Actions run through a persistent worker are not covered, since that + /// process outlives the action and serves others. With either axis on, + /// the worker reads an action's platform properties from the + /// scheduler's `StartExecute` rather than the client's request, so a + /// reservation the scheduler placed the action by is the one enforced + /// and exported to the environment. + #[serde(default)] + pub resource_enforcement: Option, + /// The command to execute on every execution request. This will be parsed as /// a command + arguments (not shell). /// Example: "run.sh" and a job with command: "sleep 5" will result in a @@ -1065,6 +1383,51 @@ pub struct LocalWorkerConfig { /// handle them. pub experimental_precondition_script: Option, + /// Milliseconds the precondition script may run before it counts as + /// failed and the action is refused as backpressure. A script that + /// hangs used to hold the action forever. + /// Default: 30000 + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub precondition_timeout_ms: u64, + + /// Bytes of an action's stdout or stderr kept in memory. Past this the + /// output spills to a file under the action directory and is uploaded + /// from disk, so an action that prints gigabytes does not grow the + /// worker by that much. 0 keeps everything in memory. + /// Default: 0 + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub max_captured_output_bytes: u64, + + /// Seconds between sweeps of `work_directory` for action directories no + /// running action owns, left by a cleanup that failed. 0 disables the + /// sweep; the directory is still purged at startup. + /// Default: 0 + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub orphan_sweep_interval_s: u64, + + /// Milliseconds between the SIGTERM a timed-out or cancelled action + /// receives and the SIGKILL that follows if it is still running, so a + /// tool that handles SIGTERM can write its own cleanup. The memory + /// reservation kill skips the grace and sends SIGKILL at once. Linux + /// only; elsewhere the kill is immediate. 0 takes the default. + /// Default: 5000 + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub kill_grace_ms: u64, + + /// Give every action its own `TMPDIR` under its action directory, + /// removed with it, so concurrent actions writing the same file name + /// under `$TMPDIR` no longer collide. A `TMPDIR` in the action's own + /// environment still wins. Code that hard-codes `/tmp` is unaffected. + /// Default: true + #[serde(default = "default_true")] + pub set_tmpdir: bool, + + /// Bazel persistent workers: how an action whose platform properties + /// carry `supports-workers=1` is run and how the pool of worker + /// processes is sized. Unset takes every default below. + #[serde(default)] + pub persistent_workers: Option, + /// Underlying CAS store that the worker will use to download CAS artifacts. /// This store must be a `FastSlowStore`. The `fast` store must be a /// `FileSystemStore` because it will use hardlinks when building out the files @@ -1091,6 +1454,15 @@ pub struct LocalWorkerConfig { /// worker. pub platform_properties: HashMap, + /// Derive the CPU and memory properties from the worker's own cgroup + /// limits instead of the values above; see `CapacityConfig`. When the + /// cgroup cannot be read (not Linux, cgroup v1, no permission) the + /// values above stand and a warning says so; if they carry neither + /// property the worker refuses to start rather than register unable to + /// take any action that asks for CPU or memory. + #[serde(default)] + pub capacity: Option, + /// An optional mapping of environment names to set for the execution /// as well as those specified in the action itself. If set, will set each /// key as an environment variable before executing the job with the value @@ -1140,6 +1512,34 @@ pub struct LocalWorkerConfig { /// error. /// Default: False. pub use_mount_namespace: Option, + + /// Whether to give each action a private `/tmp` inside its mount + /// namespace: the action's own `tmp` directory (the one `TMPDIR` points + /// at) is bound over `/tmp`, so tools that hardcode `/tmp` write there + /// too. Concurrent actions no longer collide on predictable paths under + /// `/tmp`, nothing leaks between actions through it, and what an action + /// writes there is removed with the action. + /// + /// This is only available on Linux and requires `use_namespaces` and + /// `use_mount_namespace` to be true. If explicitly set to true without + /// both the worker will exit with an error. Set it to false to keep the + /// host's `/tmp` visible to actions, for example when a tool they need + /// lives there. The bound `/tmp` carries the work volume's mount options + /// (`noexec`, `nodev`, `nosuid` if the volume has them), not the host + /// `/tmp`'s. + /// + /// The private `/tmp` lives on the worker's disk under the work + /// directory, not in memory, so it is bounded by the volume that holds + /// the work directory and seen by the disk guard, and a tool filling it + /// cannot take the worker's memory with it. + /// + /// If `/tmp` does not exist on the worker, for example in a minimal + /// container image, there is nothing for actions to collide on, so the + /// default is False there and a warning is logged at startup. + /// + /// Default: True when `use_mount_namespace` is true and `/tmp` exists, + /// otherwise False. + pub isolate_tmp: Option, } #[derive(Deserialize, Serialize, Debug, Clone)] diff --git a/nativelink-config/src/schedulers.rs b/nativelink-config/src/schedulers.rs index 4fd546c45cd..589e2eb71dd 100644 --- a/nativelink-config/src/schedulers.rs +++ b/nativelink-config/src/schedulers.rs @@ -76,6 +76,22 @@ pub enum WorkerAllocationStrategy { LeastRecentlyUsed, /// Prefer workers that have been most recently used to run a job. MostRecentlyUsed, + /// Prefer the worker running the fewest actions, least recently used + /// among equals. A worker that just joined has nothing running, so it + /// takes the next action instead of waiting for every older worker to + /// be used once more, and a burst spreads across the fleet instead of + /// filling one worker to its concurrency cap. + LeastLoaded, + /// Prefer the worker the action fits most tightly: among the workers + /// that can take it, the one whose remaining `minimum` properties + /// (memory, CPU) would be smallest afterwards, as a share of what it + /// advertises. Big actions are kept for the workers with big room and + /// small ones fill the gaps, so a fleet of mixed sizes wastes the + /// least. A `priority` property on the action that a worker carries + /// with the same value prefers that worker before the fit is judged, + /// which is what such a property is for (a zone, a cache locality). + /// Ties go to the fewest running actions, then least recently used. + BestFit, } // defaults to every 10s @@ -179,6 +195,18 @@ pub struct SimpleSpec { #[serde(default, deserialize_with = "convert_duration_with_shellexpand")] pub unacknowledged_kill_timeout_s: u64, + /// Requeue a dispatched action the worker has not acknowledged within + /// this many seconds, untried and without counting an attempt, and + /// pause the worker. The clock is the worker's own messages: a dispatch + /// is dated by the worker's last message before it, and the check runs + /// on each later message, so keep this above a few keepalive intervals. + /// Workers from v1.7.3 acknowledge every dispatch; + /// older workers never do, so leave this unset while any are + /// connected or their every action would be requeued. + /// Default: unset (off) + #[serde(default, deserialize_with = "convert_duration_with_shellexpand")] + pub dispatch_ack_timeout_s: u64, + /// Fail a queued action once no connected worker has been able to run /// it for this many seconds, instead of leaving it queued forever. An /// action counts as impossible to run when no connected worker could @@ -246,16 +274,28 @@ pub struct SimpleSpec { #[serde(default, deserialize_with = "convert_duration_with_shellexpand")] pub no_worker_action_timeout_s: u64, - /// If a job returns an internal error, times out, or loses its worker - /// this many times the scheduler completes it with `FailedPrecondition` - /// carrying the last error, a code clients do not retry. This is to help - /// prevent one rogue job from infinitely retrying and taking up a lot of - /// resources when the task itself is the one causing the server to go - /// into a bad state. + /// If a job returns an internal error or times out this many times the + /// scheduler completes it with `FailedPrecondition` carrying the last + /// error, a code clients do not retry. This is to help prevent one rogue + /// job from infinitely retrying and taking up a lot of resources when the + /// task itself is the one causing the server to go into a bad state. + /// A lost worker and a memory escalation are not the action's failures + /// and have their own budgets: `max_worker_loss_retries` and + /// `memory_escalation.max_steps`. /// Default: 3 #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] pub max_job_retries: usize, + /// How many times an action whose worker was lost (disconnected, timed + /// out, evicted, or OOM-killed as a whole) is queued again before the + /// scheduler completes it with `FailedPrecondition`. A lost worker is + /// usually not the action's fault, so these do not spend + /// `max_job_retries`; the cap only stops an action that takes a worker + /// down every time it runs. + /// Default: 10 + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub max_worker_loss_retries: usize, + /// The strategy used to assign workers jobs. #[serde(default)] pub allocation_strategy: WorkerAllocationStrategy, @@ -271,6 +311,18 @@ pub struct SimpleSpec { #[serde(default)] pub live_memory_veto: Option, + /// When a worker reports an action killed for memory (`KILLED_MEMORY`, + /// by its own reservation enforcement or by the kernel), requeue the + /// action with a larger reservation instead of failing it. Escalations + /// have their own budget (`max_steps`) and do not spend + /// `max_job_retries`. The reservation grows to the largest memory any + /// connected worker advertises (or `max_kb`); the last step reserves + /// that worker whole, memory and CPU, so the action runs alone, and only + /// a kill there fails the action: nothing in the fleet could run it. + /// Default: unset (a memory kill fails the action) + #[serde(default)] + pub memory_escalation: Option, + /// The storage backend to use for the scheduler. /// Default: memory pub experimental_backend: Option, @@ -425,6 +477,120 @@ fn default_historical_resource_memory_property_name() -> String { "memory_kb".to_string() } +#[derive(Deserialize, Serialize, Debug, Clone)] +#[serde(deny_unknown_fields)] +#[cfg_attr(feature = "dev-schema", derive(JsonSchema))] +pub struct MemoryEscalationSpec { + /// The `minimum` platform property carrying the memory reservation. + /// Default: `memory_kb` + #[serde(default = "default_historical_resource_memory_property_name")] + pub property: String, + + /// Each kill scales the reservation by this factor, in hundredths: + /// 200 doubles it, 150 adds half. + /// Default: 200 + #[serde(default = "default_memory_escalation_percent")] + pub percent: u64, + + /// Never reserve more than this many KiB; 0 takes the largest memory + /// any connected worker advertises. + /// Default: 0 + #[serde(default)] + pub max_kb: u64, + + /// The memory values of the fleet's size classes, ascending, in KiB. + /// When set, a kill steps the reservation to the next class above it + /// instead of scaling by `percent`; past the top class the last step + /// reserves the largest worker whole. The chart fills this from the + /// same classes the `historical_resource` scheduler uses. + /// Default: empty (scale by `percent`) + #[serde(default)] + pub ladder_kb: Vec, + + /// The `minimum` platform property carrying the CPU reservation. The + /// last escalation reserves the largest worker's whole memory and, + /// through this property, its whole CPU, so nothing shares the worker + /// with the action. Empty leaves CPU alone. + /// Default: `cpu_count` + #[serde(default = "default_memory_escalation_cpu_property")] + pub cpu_property: String, + + /// The most escalations one action gets before the scheduler completes + /// it with `FailedPrecondition`; 0 lets the ceiling alone bound them. + /// Default: 8 + #[serde(default = "default_memory_escalation_max_steps")] + pub max_steps: u64, +} + +fn default_memory_escalation_cpu_property() -> String { + "cpu_count".to_string() +} + +const fn default_memory_escalation_max_steps() -> u64 { + 8 +} + +const fn default_memory_escalation_percent() -> u64 { + 200 +} + +#[derive(Deserialize, Serialize, Debug, Default, Clone, Copy)] +#[serde(deny_unknown_fields)] +#[cfg_attr(feature = "dev-schema", derive(JsonSchema))] +pub struct ColdStartSpec { + /// Value for `cpu_property_name`, on the scale the workers advertise. + /// 0 leaves the property alone. + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub cpu_count: u64, + + /// Value for `memory_property_name`, in KiB. 0 leaves the property alone. + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub memory_kb: u64, + + /// Value for `disk_property_name`, in KiB. 0 leaves the property alone. + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub disk_kb: u64, +} + +fn default_disk_property_name_for_hints() -> String { + "disk_kb".to_string() +} + +/// A named point on the fleet's size ladder. List them ascending; each +/// should dominate the one before on every dimension it names. A zero +/// leaves that dimension alone. +#[derive(Deserialize, Serialize, Debug, Clone)] +#[serde(deny_unknown_fields)] +#[cfg_attr(feature = "dev-schema", derive(JsonSchema))] +pub struct SizeClassSpec { + /// Free text, referenced by hints and by the cold-start policy. + pub name: String, + + /// Value for `cpu_property_name`, on the scale the workers advertise. + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub cpu_count: u64, + + /// Value for `memory_property_name`, in KiB. + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub memory_kb: u64, + + /// Value for `disk_property_name`, in KiB. + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub disk_kb: u64, +} + +/// A class for actions whose timeout is at least `min_timeout_s`. Bazel's +/// test sizes arrive as timeouts (short 60 s, moderate 300 s, long 900 s, +/// eternal 3600 s), so this maps them to classes without a CAS fetch. +#[derive(Deserialize, Serialize, Debug, Clone)] +#[serde(deny_unknown_fields)] +#[cfg_attr(feature = "dev-schema", derive(JsonSchema))] +pub struct TimeoutClassSpec { + #[serde(deserialize_with = "convert_numeric_with_shellexpand")] + pub min_timeout_s: u64, + pub class: String, +} + #[derive(Deserialize, Serialize, Debug)] #[serde(deny_unknown_fields)] #[cfg_attr(feature = "dev-schema", derive(JsonSchema))] @@ -453,7 +619,10 @@ pub struct HistoricalResourceSpec { )] pub refresh_interval_s: u64, - /// Platform property name used for CPU minimums. + /// Platform property name used for CPU minimums. The nested scheduler + /// must declare it as a `minimum` property (as it must the memory and + /// disk names): an undeclared property is matched as an exact string, + /// and an action carrying a number would then match no worker. /// Default: `cpu_count` #[serde( default = "default_historical_resource_cpu_property_name", @@ -462,6 +631,7 @@ pub struct HistoricalResourceSpec { pub cpu_property_name: String, /// Platform property name used for memory minimums, expressed in KiB. + /// Declared as `minimum` on the nested scheduler, like the CPU name. /// Default: `memory_kb` #[serde( default = "default_historical_resource_memory_property_name", @@ -469,6 +639,45 @@ pub struct HistoricalResourceSpec { )] pub memory_property_name: String, + /// Platform property name used for disk minimums, expressed in KiB. + /// Declared as `minimum` on the nested scheduler, like the CPU name. + /// Default: `disk_kb` + #[serde( + default = "default_disk_property_name_for_hints", + deserialize_with = "convert_string_with_shellexpand" + )] + pub disk_property_name: String, + + /// The fleet's size ladder. A hint may name a `class` instead of + /// numbers, and the cold-start policy below picks one for actions no + /// hint matches. Empty means raw numbers only. + #[serde(default)] + pub classes: Vec, + + /// The class an untagged action gets when no hint and no rule below + /// matches. Takes precedence over `cold_start` when both are set. + #[serde(default)] + pub default_class: Option, + + /// Class by Bazel `action_mnemonic` for untagged actions; a known heavy + /// mnemonic (`TestRunner`, `Link`) starts in the right class. + #[serde(default)] + pub class_by_mnemonic: HashMap, + + /// Class by the action's timeout, the largest `min_timeout_s` at or + /// below the timeout wins. Checked before `default_class`, after the + /// mnemonic rule. + #[serde(default)] + pub class_by_timeout: Vec, + + /// Reservation given to an action that no hint matches and the client + /// left untagged. Only a property that is absent is filled; a value the + /// client sent is kept. Without this an untagged action costs the + /// scheduler's ledger nothing, so any number of them can land on one + /// worker. + #[serde(default)] + pub cold_start: Option, + /// The nested scheduler to use after applying resource hints. pub scheduler: Box, } diff --git a/nativelink-config/src/stores.rs b/nativelink-config/src/stores.rs index d0adab31308..7d312f5667d 100644 --- a/nativelink-config/src/stores.rs +++ b/nativelink-config/src/stores.rs @@ -1414,6 +1414,12 @@ pub enum StoreType { Cas, /// The store is an action cache. Ac, + /// A write-only event sink using `ByteStream.Write`. Only string keys and + /// exact upload sizes are accepted. The resource hash identifies the key, + /// not the payload; the size describes the payload. The upstream must be an + /// event receiver, not a CAS. Reads and existence checks are unsupported. + /// Use this for BEP/origin event forwarding, never for cache data. + EventSink, } #[derive(Serialize, Deserialize, Debug, Clone)] diff --git a/nativelink-proto/com/github/trace_machina/nativelink/remote_execution/worker_api.proto b/nativelink-proto/com/github/trace_machina/nativelink/remote_execution/worker_api.proto index 6cd3840d970..35c36987705 100644 --- a/nativelink-proto/com/github/trace_machina/nativelink/remote_execution/worker_api.proto +++ b/nativelink-proto/com/github/trace_machina/nativelink/remote_execution/worker_api.proto @@ -62,7 +62,12 @@ message KeepAliveRequest { /// Request object for going away requests. message GoingAwayRequest { - reserved 1; // NextId. + reserved 1; + /// The worker is shutting down but still running actions: stop + /// dispatching to it and let those finish. It leaves the pool when its + /// stream closes. Without this the scheduler removes the worker at + /// once and requeues everything it holds. + bool drain = 2; // NextId: 3. } /// Represents the initial request sent to the scheduler informing the @@ -125,6 +130,32 @@ message ExecuteComplete { string operation_id = 1; } +/// How an action ended, as the worker saw it. Separate from the exit code: +/// a kill by the worker or the kernel is a sizing fact, not a tool failure. +enum ResourceOutcome { + RESOURCE_OUTCOME_UNKNOWN = 0; + /// The process exited on its own, whatever its exit code. + COMPLETED = 1; + /// Killed for memory: by the worker's reservation enforcement, or by + /// the kernel with the last sample within 10% of the memory limit. + KILLED_MEMORY = 2; + /// Killed for disk. Reserved for disk enforcement. + KILLED_DISK = 3; + /// Killed by the action timeout. + KILLED_TIMEOUT = 4; + /// Killed from outside: a scheduler cancel, or a signal nobody here sent. + KILLED_EXTERNAL = 5; +} + +/// The reservation an action was admitted under, echoed by the worker so +/// the measurement and the number it was measured against travel together. +message Reservation { + uint64 cpu_count = 1; + uint64 memory_kb = 2; + uint64 disk_kb = 3; + reserved 4; // NextId. +} + /// Resource usage observed by the worker while running one action. message ActionResourceUsage { /// Peak resident memory observed for the action process tree. @@ -144,7 +175,25 @@ message ActionResourceUsage { /// action's process group. uint64 cpu_time_ms = 5; - reserved 6; // NextId. + /// Wall time from spawn to exit, in milliseconds; cores used is + /// cpu_time_ms / wall_time_ms. + uint64 wall_time_ms = 6; + + /// Peak bytes under the action's work directory, in KiB. Not yet + /// measured; zero. + uint64 peak_disk_kb = 7; + + /// How the action ended. + ResourceOutcome outcome = 8; + + /// True when a worker-side limit, not the kernel or a timeout, ended + /// the action. + bool enforced = 9; + + /// What the action was admitted with, from its platform properties. + Reservation reserved = 10; + + reserved 11; // NextId. } /// Result sent back from the server when a node connects. @@ -152,9 +201,58 @@ message ConnectionResult { /// The worker ID to place in the action results generated by this worker. string worker_id = 1; + /// The scheduler understands ExecuteAccepted and ExecuteDeclined. A + /// worker sends them only to a scheduler that says so: an older + /// scheduler would drop them, and a declined action would then wait on + /// a worker that never runs it. + bool dispatch_ack = 2; + + /// The platform property the scheduler reads an action's memory + /// reservation from, when it vetoes placement on the worker's reported + /// free memory (`live_memory_veto`). The worker declines a dispatch for + /// load against the same property, so the two sides cannot disagree. + /// Empty when the scheduler does not veto; the worker then never + /// declines for load. + string memory_property = 3; + + reserved 4; // NextId. +} + +/// The worker took the action it was sent and is about to run it. Until +/// this arrives the scheduler has only charged its ledger on the send. +message ExecuteAccepted { + string operation_id = 1; reserved 2; // NextId. } +/// The worker will not run the action it was sent. The scheduler restores +/// its ledger and requeues the action without counting an attempt, and +/// pauses the worker until it reports room again. +message ExecuteDeclined { + string operation_id = 1; + + enum Reason { + REASON_UNSPECIFIED = 0; + /// The worker already has `max_inflight_tasks` actions. + AT_CAPACITY = 1; + /// The action's memory reservation is more than the worker has free. + LOAD = 2; + /// The worker is shutting down, or is single-use and already took + /// its action. + SHUTTING_DOWN = 3; + } + Reason reason = 2; + + /// Free text for the log. + string detail = 3; + + /// For LOAD: what the action asked for and what the worker had, in KiB. + uint64 needed_kb = 4; + uint64 free_kb = 5; + + reserved 6; // NextId. +} + /// Request to kill a running operation sent from the scheduler to a worker. message KillOperationRequest { /// The the operation id for the operation to be killed. @@ -221,8 +319,14 @@ message UpdateForScheduler { /// Notify that the execution has completed, but result is uploading. ExecuteComplete execute_complete = 5; + + /// The worker took the action. + ExecuteAccepted execute_accepted = 6; + + /// The worker will not run the action. + ExecuteDeclined execute_declined = 7; } - reserved 6; // NextId. + reserved 8; // NextId. } message StartExecute { diff --git a/nativelink-proto/genproto/com.github.trace_machina.nativelink.remote_execution.pb.rs b/nativelink-proto/genproto/com.github.trace_machina.nativelink.remote_execution.pb.rs index 2f4954e6a60..6187780b507 100644 --- a/nativelink-proto/genproto/com.github.trace_machina.nativelink.remote_execution.pb.rs +++ b/nativelink-proto/genproto/com.github.trace_machina.nativelink.remote_execution.pb.rs @@ -32,7 +32,16 @@ pub struct KeepAliveRequest { } /// / Request object for going away requests. #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] -pub struct GoingAwayRequest {} +pub struct GoingAwayRequest { + /// / The worker is shutting down but still running actions: stop + /// / dispatching to it and let those finish. It leaves the pool when its + /// / stream closes. Without this the scheduler removes the worker at + /// / once and requeues everything it holds. + /// + /// NextId: 3. + #[prost(bool, tag = "2")] + pub drain: bool, +} /// / Represents the initial request sent to the scheduler informing the /// / scheduler about this worker's capabilities and metadata. #[derive(Clone, PartialEq, ::prost::Message)] @@ -103,6 +112,17 @@ pub struct ExecuteComplete { #[prost(string, tag = "1")] pub operation_id: ::prost::alloc::string::String, } +/// / The reservation an action was admitted under, echoed by the worker so +/// / the measurement and the number it was measured against travel together. +#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] +pub struct Reservation { + #[prost(uint64, tag = "1")] + pub cpu_count: u64, + #[prost(uint64, tag = "2")] + pub memory_kb: u64, + #[prost(uint64, tag = "3")] + pub disk_kb: u64, +} /// / Resource usage observed by the worker while running one action. #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct ActionResourceUsage { @@ -123,6 +143,24 @@ pub struct ActionResourceUsage { /// / action's process group. #[prost(uint64, tag = "5")] pub cpu_time_ms: u64, + /// / Wall time from spawn to exit, in milliseconds; cores used is + /// / cpu_time_ms / wall_time_ms. + #[prost(uint64, tag = "6")] + pub wall_time_ms: u64, + /// / Peak bytes under the action's work directory, in KiB. Not yet + /// / measured; zero. + #[prost(uint64, tag = "7")] + pub peak_disk_kb: u64, + /// / How the action ended. + #[prost(enumeration = "ResourceOutcome", tag = "8")] + pub outcome: i32, + /// / True when a worker-side limit, not the kernel or a timeout, ended + /// / the action. + #[prost(bool, tag = "9")] + pub enforced: bool, + /// / What the action was admitted with, from its platform properties. + #[prost(message, optional, tag = "10")] + pub reserved: ::core::option::Option, } /// / Result sent back from the server when a node connects. #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] @@ -130,6 +168,94 @@ pub struct ConnectionResult { /// / The worker ID to place in the action results generated by this worker. #[prost(string, tag = "1")] pub worker_id: ::prost::alloc::string::String, + /// / The scheduler understands ExecuteAccepted and ExecuteDeclined. A + /// / worker sends them only to a scheduler that says so: an older + /// / scheduler would drop them, and a declined action would then wait on + /// / a worker that never runs it. + #[prost(bool, tag = "2")] + pub dispatch_ack: bool, + /// / The platform property the scheduler reads an action's memory + /// / reservation from, when it vetoes placement on the worker's reported + /// / free memory (`live_memory_veto`). The worker declines a dispatch for + /// / load against the same property, so the two sides cannot disagree. + /// / Empty when the scheduler does not veto; the worker then never + /// / declines for load. + #[prost(string, tag = "3")] + pub memory_property: ::prost::alloc::string::String, +} +/// / The worker took the action it was sent and is about to run it. Until +/// / this arrives the scheduler has only charged its ledger on the send. +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct ExecuteAccepted { + #[prost(string, tag = "1")] + pub operation_id: ::prost::alloc::string::String, +} +/// / The worker will not run the action it was sent. The scheduler restores +/// / its ledger and requeues the action without counting an attempt, and +/// / pauses the worker until it reports room again. +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct ExecuteDeclined { + #[prost(string, tag = "1")] + pub operation_id: ::prost::alloc::string::String, + #[prost(enumeration = "execute_declined::Reason", tag = "2")] + pub reason: i32, + /// / Free text for the log. + #[prost(string, tag = "3")] + pub detail: ::prost::alloc::string::String, + /// / For LOAD: what the action asked for and what the worker had, in KiB. + #[prost(uint64, tag = "4")] + pub needed_kb: u64, + #[prost(uint64, tag = "5")] + pub free_kb: u64, +} +/// Nested message and enum types in `ExecuteDeclined`. +pub mod execute_declined { + #[derive( + Clone, + Copy, + Debug, + PartialEq, + Eq, + Hash, + PartialOrd, + Ord, + ::prost::Enumeration + )] + #[repr(i32)] + pub enum Reason { + Unspecified = 0, + /// / The worker already has `max_inflight_tasks` actions. + AtCapacity = 1, + /// / The action's memory reservation is more than the worker has free. + Load = 2, + /// / The worker is shutting down, or is single-use and already took + /// / its action. + ShuttingDown = 3, + } + impl Reason { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + Self::Unspecified => "REASON_UNSPECIFIED", + Self::AtCapacity => "AT_CAPACITY", + Self::Load => "LOAD", + Self::ShuttingDown => "SHUTTING_DOWN", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "REASON_UNSPECIFIED" => Some(Self::Unspecified), + "AT_CAPACITY" => Some(Self::AtCapacity), + "LOAD" => Some(Self::Load), + "SHUTTING_DOWN" => Some(Self::ShuttingDown), + _ => None, + } + } + } } /// / Request to kill a running operation sent from the scheduler to a worker. #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] @@ -175,7 +301,7 @@ pub mod update_for_worker { /// / Communication from the worker to the scheduler. #[derive(Clone, PartialEq, ::prost::Message)] pub struct UpdateForScheduler { - #[prost(oneof = "update_for_scheduler::Update", tags = "1, 2, 3, 4, 5")] + #[prost(oneof = "update_for_scheduler::Update", tags = "1, 2, 3, 4, 5, 6, 7")] pub update: ::core::option::Option, } /// Nested message and enum types in `UpdateForScheduler`. @@ -212,6 +338,12 @@ pub mod update_for_scheduler { /// / Notify that the execution has completed, but result is uploading. #[prost(message, tag = "5")] ExecuteComplete(super::ExecuteComplete), + /// / The worker took the action. + #[prost(message, tag = "6")] + ExecuteAccepted(super::ExecuteAccepted), + /// / The worker will not run the action. + #[prost(message, tag = "7")] + ExecuteDeclined(super::ExecuteDeclined), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -259,6 +391,52 @@ pub struct HistoricalExecuteResponse { super::super::super::super::super::build::bazel::remote::execution::v2::ExecuteResponse, >, } +/// / How an action ended, as the worker saw it. Separate from the exit code: +/// / a kill by the worker or the kernel is a sizing fact, not a tool failure. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] +pub enum ResourceOutcome { + Unknown = 0, + /// / The process exited on its own, whatever its exit code. + Completed = 1, + /// / Killed for memory: by the worker's reservation enforcement, or by + /// / the kernel with the last sample within 10% of the memory limit. + KilledMemory = 2, + /// / Killed for disk. Reserved for disk enforcement. + KilledDisk = 3, + /// / Killed by the action timeout. + KilledTimeout = 4, + /// / Killed from outside: a scheduler cancel, or a signal nobody here sent. + KilledExternal = 5, +} +impl ResourceOutcome { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + Self::Unknown => "RESOURCE_OUTCOME_UNKNOWN", + Self::Completed => "COMPLETED", + Self::KilledMemory => "KILLED_MEMORY", + Self::KilledDisk => "KILLED_DISK", + Self::KilledTimeout => "KILLED_TIMEOUT", + Self::KilledExternal => "KILLED_EXTERNAL", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "RESOURCE_OUTCOME_UNKNOWN" => Some(Self::Unknown), + "COMPLETED" => Some(Self::Completed), + "KILLED_MEMORY" => Some(Self::KilledMemory), + "KILLED_DISK" => Some(Self::KilledDisk), + "KILLED_TIMEOUT" => Some(Self::KilledTimeout), + "KILLED_EXTERNAL" => Some(Self::KilledExternal), + _ => None, + } + } +} /// Generated client implementations. pub mod worker_api_client { #![allow( diff --git a/nativelink-scheduler/BUILD.bazel b/nativelink-scheduler/BUILD.bazel index 2601e67f86b..8f2f4f632d8 100644 --- a/nativelink-scheduler/BUILD.bazel +++ b/nativelink-scheduler/BUILD.bazel @@ -10,6 +10,7 @@ load( rust_library( name = "nativelink-scheduler", srcs = [ + "src/admin.rs", "src/api_worker_scheduler.rs", "src/awaited_action_db/awaited_action.rs", "src/awaited_action_db/mod.rs", @@ -70,12 +71,16 @@ rust_test_suite( timeout = "short", srcs = [ "tests/action_messages_test.rs", + "tests/admin_test.rs", "tests/backpressure_test.rs", "tests/cache_lookup_scheduler_test.rs", + "tests/dispatch_ack_test.rs", + "tests/escalation_test.rs", "tests/execution_active_count_test.rs", "tests/fleet_capabilities_test.rs", "tests/historical_resource_scheduler_test.rs", "tests/live_memory_veto_test.rs", + "tests/placement_test.rs", "tests/platform_property_manager_test.rs", "tests/property_modifier_scheduler_test.rs", "tests/queue_order_test.rs", diff --git a/nativelink-scheduler/src/admin.rs b/nativelink-scheduler/src/admin.rs new file mode 100644 index 00000000000..115b5738c06 --- /dev/null +++ b/nativelink-scheduler/src/admin.rs @@ -0,0 +1,142 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Read-only views for the admin API: what the scheduler has queued and what +//! it has connected. A provisioner sizes pods from the first and drains +//! through the second, instead of reading the scheduler's store behind its +//! back. + +use std::collections::HashMap; +use std::time::{SystemTime, UNIX_EPOCH}; + +use futures::StreamExt; +use nativelink_error::{Error, ResultExt}; +use nativelink_util::action_messages::ActionStage; +use nativelink_util::metrics::record_awaited_action_orphan; +use nativelink_util::operation_state_manager::{ + ClientStateManager, OperationFilter, OperationStageFlags, OrderDirection, +}; +use serde::{Deserialize, Serialize}; +use tracing::warn; + +use crate::simple_scheduler_state_manager::is_lost_record; +use crate::worker_scheduler::WorkerSummary; + +/// A queued action as the admin API lists it: enough to size a pod for it +/// and to tell how long it has waited. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct QueuedDemand { + pub operation_id: String, + /// Milliseconds since the epoch when the action was queued. + pub queued_since_ms: u64, + pub priority: i32, + /// The action's properties as the scheduler matches them, reservations + /// included. + pub platform_properties: HashMap, + /// Milliseconds since the epoch when a client last checked in on the + /// operation. A queue entry no client has touched for longer than the + /// client timeout is one the scheduler will not dispatch. + pub client_last_seen_ms: Option, +} + +/// Every queued action of the scheduler, oldest first as the store yields +/// them. +pub async fn queued_demand(scheduler: &dyn ClientStateManager) -> Result, Error> { + let mut stream = scheduler + .filter_operations(OperationFilter { + stages: OperationStageFlags::Queued, + // The Redis backend serves the queue in one direction only, the + // matcher's: highest priority first, then oldest. + order_by_priority_direction: Some(OrderDirection::Desc), + ..Default::default() + }) + .await + .err_tip(|| "Listing queued operations for the admin API")?; + let mut out = Vec::new(); + while let Some(result) = stream.next().await { + // The filter read the record once; each read below is a fresh one, + // and the record can have moved on or gone in between. One that is + // gone or cannot be decoded is left out and counted, one that has + // since been dispatched is no longer demand, and a store failure is + // the listing's error. + let Some((state, _)) = still_listed(result.as_state().await)? else { + continue; + }; + if !matches!(state.stage, ActionStage::Queued) { + continue; + } + let Some((action_info, _)) = still_listed(result.as_action_info().await)? else { + continue; + }; + let Some(client_last_seen) = still_listed(result.client_last_seen().await)? else { + continue; + }; + let client_last_seen_ms = client_last_seen.map(epoch_ms); + out.push(QueuedDemand { + operation_id: state.client_operation_id.to_string(), + queued_since_ms: epoch_ms(action_info.insert_timestamp), + priority: action_info.priority, + platform_properties: action_info.platform_properties.clone(), + client_last_seen_ms, + }); + } + Ok(out) +} + +/// A read of a listed record: the value, `None` for a record that is gone +/// or cannot be decoded (skipped and counted), or the store's error. +fn still_listed(read: Result) -> Result, Error> { + match read { + Ok(value) => Ok(Some(value)), + Err(err) if is_lost_record(&err) => { + warn!( + ?err, + "Queued operation listed but its record cannot be read; left out of the demand" + ); + record_awaited_action_orphan("admin"); + Ok(None) + } + Err(err) => Err(err).err_tip(|| "Reading a queued operation for the admin API"), + } +} + +fn epoch_ms(at: SystemTime) -> u64 { + u64::try_from( + at.duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis(), + ) + .unwrap_or(u64::MAX) +} + +/// The demand listing as the JSON body the admin API serves. +pub async fn queued_demand_json(scheduler: &dyn ClientStateManager) -> Result { + let demand = queued_demand(scheduler).await?; + serde_json::to_string(&demand).map_err(|e| { + Error::new( + nativelink_error::Code::Internal, + format!("Serializing the demand listing: {e}"), + ) + }) +} + +/// The worker listing as the JSON body the admin API serves. +pub fn workers_json(workers: &[WorkerSummary]) -> Result { + serde_json::to_string(workers).map_err(|e| { + Error::new( + nativelink_error::Code::Internal, + format!("Serializing the worker listing: {e}"), + ) + }) +} diff --git a/nativelink-scheduler/src/api_worker_scheduler.rs b/nativelink-scheduler/src/api_worker_scheduler.rs index d3565ff1446..7392fddf899 100644 --- a/nativelink-scheduler/src/api_worker_scheduler.rs +++ b/nativelink-scheduler/src/api_worker_scheduler.rs @@ -22,7 +22,7 @@ use std::time::{Instant, SystemTime, UNIX_EPOCH}; use async_lock::Mutex; use futures::{StreamExt, future}; use lru::LruCache; -use nativelink_config::schedulers::WorkerAllocationStrategy; +use nativelink_config::schedulers::{MemoryEscalationSpec, WorkerAllocationStrategy}; use nativelink_error::{Code, Error, ResultExt, error_if, make_err, make_input_err}; use nativelink_metric::{ MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent, @@ -32,15 +32,17 @@ use nativelink_proto::com::github::trace_machina::nativelink::events::{ Event, OriginEvent, ResponseEvent, event, response_event, }; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{ - ActionResourceUsage, WorkerLoad, + ActionResourceUsage, ResourceOutcome, WorkerLoad, }; -use nativelink_util::action_messages::{OperationId, WorkerId}; +use nativelink_util::action_messages::{ActionStage, OperationId, WorkerId}; use nativelink_util::metrics::{ - WorkerDisconnectReason, record_execution_cpu_time, record_execution_peak_memory, - record_worker_connected, record_worker_disconnected, record_worker_keepalive, - record_worker_keepalive_gap, record_worker_state, + WorkerDisconnectReason, record_dispatch_requeue, record_execution_cpu_time, + record_execution_peak_memory, record_worker_connected, record_worker_disconnected, + record_worker_keepalive, record_worker_keepalive_gap, record_worker_state, +}; +use nativelink_util::operation_state_manager::{ + Decline, Escalation, UpdateOperationType, WorkerStateManager, }; -use nativelink_util::operation_state_manager::{UpdateOperationType, WorkerStateManager}; use nativelink_util::origin_event::get_node_id; use nativelink_util::platform_properties::{PlatformProperties, PlatformPropertyValue}; use nativelink_util::shutdown_guard::ShutdownGuard; @@ -52,6 +54,49 @@ use tracing::{debug, error, info, trace, warn}; /// at once while checking which running operations were revoked. const MAX_CONCURRENT_REVOKED_CHECKS: usize = 32; +/// How many of the action's `priority` properties the worker carries with +/// the same value. A priority property never restricts placement; it says +/// where the action would rather be. +fn priority_matches(action: &PlatformProperties, worker: &Worker) -> usize { + action + .properties + .iter() + .filter(|(name, value)| { + matches!(value, PlatformPropertyValue::Priority(_)) + && worker.platform_properties.properties.get(*name) == Some(*value) + }) + .count() +} + +/// How much of the worker's advertised capacity would be left after this +/// action, summed over the action's `minimum` properties, in thousandths. +/// Lower is a tighter fit. An action with no numbers fits every worker +/// equally, and a dimension the worker does not advertise is ignored. +fn fit_leftover_permille(action: &PlatformProperties, worker: &Worker) -> u64 { + let mut leftover = 0u64; + for (name, value) in &action.properties { + let PlatformPropertyValue::Minimum(needed) = value else { + continue; + }; + let ( + Some(PlatformPropertyValue::Minimum(available)), + Some(PlatformPropertyValue::Minimum(total)), + ) = ( + worker.platform_properties.properties.get(name), + worker.total_platform_properties.properties.get(name), + ) + else { + continue; + }; + if *total == 0 { + continue; + } + let remaining = available.saturating_sub(*needed); + leftover += remaining.saturating_mul(1000) / total; + } + leftover +} + /// How many property shapes `static_verdicts` holds before it is emptied. const MAX_STATIC_VERDICTS: usize = 4096; @@ -80,6 +125,12 @@ pub struct SchedulerMetrics { pub keep_alive_updates: AtomicU64, /// Total number of worker timeouts. pub worker_timeouts: AtomicU64, + /// Dispatches the worker declined, requeued untried. + pub dispatches_declined: AtomicU64, + /// Dispatches that found the worker's channel full, requeued untried. + pub dispatch_channel_full: AtomicU64, + /// Dispatches requeued because the worker never acknowledged them. + pub dispatches_unacknowledged: AtomicU64, } use crate::match_outcome::{ @@ -89,7 +140,7 @@ use crate::platform_property_manager::PlatformPropertyManager; use crate::worker::{ActionInfoWithProps, Worker, WorkerTimestamp, WorkerUpdate}; use crate::worker_capability_index::WorkerCapabilityIndex; use crate::worker_registry::SharedWorkerRegistry; -use crate::worker_scheduler::WorkerScheduler; +use crate::worker_scheduler::{WorkerScheduler, WorkerSummary}; #[derive(Debug)] struct Workers(LruCache); @@ -140,8 +191,15 @@ struct ApiWorkerSchedulerImpl { /// The `minimum` property compared against each worker's reported free /// memory before placement; unset means no veto. live_memory_veto: Option, + /// How a memory kill is turned into a retry with a larger reservation; + /// unset means the kill fails the action. + memory_escalation: Option, /// A channel to notify the matching engine that the worker pool has changed. worker_change_notify: Arc, + /// Bumped with every notification, under the same lock as placement, so + /// a matching pass can tell that room opened while it was walking the + /// queue and offer that room to the oldest waiting action first. + capacity_generation: Arc, /// Worker registry for tracking worker liveness. worker_registry: SharedWorkerRegistry, @@ -273,6 +331,7 @@ impl ApiWorkerSchedulerImpl { worker_id: &WorkerId, timestamp: WorkerTimestamp, load: Option, + keepalive: bool, ) -> Result<(), Error> { let worker = self.workers.0.peek_mut(worker_id).ok_or_else(|| { make_input_err!( @@ -288,28 +347,49 @@ impl ApiWorkerSchedulerImpl { ); record_worker_keepalive_gap(timestamp.saturating_sub(worker.last_update_timestamp)); worker.last_update_timestamp = timestamp; - // A keepalive is the worker saying it is ready to be asked again, - // so a pause from backpressure lasts one keepalive interval. The - // matcher has to hear about it, or the requeued action waits for - // the next unrelated change. - if worker.is_paused { - worker.is_paused = false; - record_worker_state("paused", false); - self.worker_change_notify.notify_one(); + // Any other message (an acknowledgement, a decline, an execute + // result) proves the worker is alive and nothing more. A decline in + // particular is not the worker saying it is ready to be asked + // again, so the pause it took stands until a real keepalive. + if !keepalive { + return Ok(()); } - // A message without a load (an older worker, or a refresh from an - // execute result) leaves the last report in place. A report with + // A keepalive without a load (an older worker) leaves the last + // report in place. A report with // more room than the last one can make a vetoed action eligible - // again, so the matcher hears about it; less room never does. - if let Some(load) = load { - let more_room = worker + // again; less room never does. + let more_room = load.is_some_and(|load| { + let more = worker .last_load .as_ref() .is_none_or(|last| load.free_memory_kb > last.free_memory_kb); worker.last_load = Some(load); - if more_room { - self.worker_change_notify.notify_one(); - } + more + }); + // A keepalive is the worker saying it is ready to be asked again, + // so a pause from backpressure lasts one keepalive interval. A + // pause taken for a decline for load is the exception: it waits + // for a keepalive whose free memory, read above, covers what the + // declined action wanted. A worker that reports no load is taken + // at its word. + let load_fits = match (worker.pause_needs_kb, worker.last_load) { + (Some(needs_kb), Some(load)) => load.free_memory_kb >= needs_kb, + _ => true, + }; + let resumed = worker.is_paused && load_fits; + if resumed { + worker.is_paused = false; + worker.pause_needs_kb = None; + record_worker_state("paused", false); + } + // Either is room the fleet did not have a moment ago, so the + // capacity generation moves and the matcher is woken, together: + // a pass in flight offers the room to its parked actions first, + // and a pass not in flight starts. `capacity_changed` by hand, + // since `worker` still borrows the map. + if resumed || more_room { + self.capacity_generation.fetch_add(1, Ordering::Relaxed); + self.worker_change_notify.notify_one(); } trace!( @@ -342,7 +422,7 @@ impl ApiWorkerSchedulerImpl { // the multi-threaded runtime works. let worker = self.workers.peek_mut(&worker_id).unwrap(); let res = worker - .send_initial_connection_result() + .send_initial_connection_result(self.live_memory_veto.as_deref()) .err_tip(|| "Failed to send initial connection result to worker"); if let Err(err) = &res { error!( @@ -361,7 +441,7 @@ impl ApiWorkerSchedulerImpl { } else { record_worker_connected(); } - self.worker_change_notify.notify_one(); + self.capacity_changed(); res } @@ -377,7 +457,7 @@ impl ApiWorkerSchedulerImpl { self.fleet_changed(); self.local_fleet_change_notify.notify_one(); - self.worker_change_notify.notify_one(); + self.capacity_changed(); result } @@ -398,6 +478,13 @@ impl ApiWorkerSchedulerImpl { self.static_verdicts.clear(); } + /// Something about the fleet's room to run actions changed: a worker + /// joined, left, paused, resumed, drained or finished an action. + fn capacity_changed(&self) { + self.capacity_generation.fetch_add(1, Ordering::Relaxed); + self.worker_change_notify.notify_one(); + } + /// The distinct properties this scheduler's workers registered with. fn fleet_shapes(&self) -> Vec { let mut seen = HashSet::new(); @@ -462,7 +549,7 @@ impl ApiWorkerSchedulerImpl { record_worker_state("draining", is_draining); } worker.is_draining = is_draining; - self.worker_change_notify.notify_one(); + self.capacity_changed(); Ok(()) } @@ -573,10 +660,19 @@ impl ApiWorkerSchedulerImpl { // The ledger only knows what actions declared; the worker's own // report of what it has left catches the ones that declared too // little. A worker that reports nothing is never vetoed. + // An action that asks for the worker's whole advertised memory + // (the last escalation step) is not vetoed: no worker ever + // reports that much free, since its own binary, the kernel and + // the page cache hold some, and the ledger already keeps it + // alone there. if let (Some(property), Some(load)) = (self.live_memory_veto.as_deref(), w.last_load) && let Some(PlatformPropertyValue::Minimum(needed_kb)) = platform_properties.properties.get(property) && *needed_kb > load.free_memory_kb + && matches!( + w.total_platform_properties.properties.get(property), + Some(PlatformPropertyValue::Minimum(advertised_kb)) if *needed_kb < *advertised_kb + ) { if full_worker_logging { info!( @@ -608,6 +704,29 @@ impl ApiWorkerSchedulerImpl { .filter(|(worker_id, _)| candidates.contains(worker_id)) .find(&worker_matches) .map(|(_, w)| w.id.clone()), + + // Fewest running actions wins; the scan runs from the least + // recently used end so `min_by_key` keeps that one on a tie. + WorkerAllocationStrategy::LeastLoaded => workers_iter + .rev() + .filter(|(worker_id, _)| candidates.contains(worker_id)) + .filter(&worker_matches) + .min_by_key(|(_, w)| w.running_action_infos.len()) + .map(|(_, w)| w.id.clone()), + + WorkerAllocationStrategy::BestFit => workers_iter + .rev() + .filter(|(worker_id, _)| candidates.contains(worker_id)) + .filter(&worker_matches) + .min_by_key(|(_, w)| { + ( + // More matching priority properties first. + usize::MAX - priority_matches(platform_properties, w), + fit_leftover_permille(platform_properties, w), + w.running_action_infos.len(), + ) + }) + .map(|(_, w)| w.id.clone()), } } @@ -676,6 +795,204 @@ impl ApiWorkerSchedulerImpl { MatchOutcome::WaitingForCapacity } + /// The largest memory reservation any connected worker could hold. + /// The largest memory any connected worker advertises under `property`, + /// and that worker's whole CPU under `cpu_property` (`None` when the + /// property is empty or the worker does not advertise it). Reserving + /// both is what runs an action alone there. + fn fleet_largest_worker(&self, property: &str, cpu_property: &str) -> (u64, Option) { + let minimum = |worker: &Worker, name: &str| match worker + .total_platform_properties + .properties + .get(name) + { + Some(PlatformPropertyValue::Minimum(value)) => Some(*value), + _ => None, + }; + // A draining worker is leaving: escalating to it would strand the + // action at a size the remaining fleet cannot serve. + self.workers + .iter() + .filter(|(_, worker)| !worker.is_draining) + .filter_map(|(_, worker)| { + minimum(worker, property).map(|memory_kb| { + let cpu = if cpu_property.is_empty() { + None + } else { + minimum(worker, cpu_property) + }; + (memory_kb, cpu) + }) + }) + .max_by_key(|(memory_kb, _)| *memory_kb) + .unwrap_or((0, None)) + } + + /// Turns a result that the worker's last usage report says was a memory + /// kill into a requeue with a larger reservation, when escalation is + /// configured and there is room to grow. Returns the update to apply + /// and whether it was escalated. + fn maybe_escalate( + &self, + worker_id: &WorkerId, + operation_id: &OperationId, + update: UpdateOperationType, + ) -> (UpdateOperationType, bool) { + let Some(policy) = &self.memory_escalation else { + return (update, false); + }; + let UpdateOperationType::UpdateWithActionStage(ActionStage::Completed(result)) = &update + else { + return (update, false); + }; + let Some(pending) = self + .workers + .peek(worker_id) + .and_then(|worker| worker.running_action_infos.get(operation_id)) + else { + return (update, false); + }; + let Some(usage) = pending.last_usage.as_ref() else { + return (update, false); + }; + if usage.outcome != ResourceOutcome::KilledMemory as i32 { + return (update, false); + } + // What it ran under: the worker's echo, else what the scheduler sent. + let reserved_kb = usage + .reserved + .map(|r| r.memory_kb) + .filter(|kb| *kb > 0) + .or_else(|| { + match pending + .action_info + .platform_properties + .properties + .get(&policy.property) + { + Some(PlatformPropertyValue::Minimum(kb)) => Some(*kb), + _ => None, + } + }) + .unwrap_or(0); + let (fleet_memory_kb, fleet_cpu) = + self.fleet_largest_worker(&policy.property, &policy.cpu_property); + // `max_kb` caps the fleet's ceiling; it cannot raise it past what a + // worker declares, or the escalation lands where nothing can run. + let ceiling_kb = if policy.max_kb > 0 { + fleet_memory_kb.min(policy.max_kb) + } else { + fleet_memory_kb + }; + if ceiling_kb == 0 { + return (update, false); + } + // The last step reserves the largest worker whole: its memory and, + // when the fleet advertises one, its CPU, so nothing shares the + // worker with the action. A kill there is the end: no worker in + // the fleet can run the action, and the client hears exactly that. + let cpu_held = match pending + .action_info + .platform_properties + .properties + .get(&policy.cpu_property) + { + Some(PlatformPropertyValue::Minimum(cores)) => *cores, + _ => 0, + }; + let whole_worker_cpu = fleet_cpu.filter(|cores| *cores > 0); + let already_whole = + reserved_kb >= ceiling_kb && whole_worker_cpu.is_none_or(|cores| cpu_held >= cores); + if already_whole { + warn!( + ?worker_id, + ?operation_id, + reserved_kb, + ceiling_kb, + "Action killed for memory with the largest worker reserved whole; no worker in the fleet can run it" + ); + let mut failed = result.clone(); + let no_worker = make_err!( + Code::FailedPrecondition, + "Action was killed for memory at {reserved_kb} KiB with the largest worker in the fleet ({ceiling_kb} KiB) reserved for it alone; no worker can run it. It needs a bigger worker or less memory" + ); + failed.error = Some(match result.error.clone() { + Some(worker_error) => no_worker.merge(worker_error), + None => no_worker, + }); + return ( + UpdateOperationType::UpdateWithActionStage(ActionStage::Completed(failed)), + false, + ); + } + let base_kb = if reserved_kb > 0 { + reserved_kb + } else { + usage.peak_memory_kb + }; + if base_kb == 0 { + // Nothing to grow from: no reservation and an unsampled kill. + // Scaling zero would hand out the smallest steps the ladder or + // the percent allow and burn the budget on them. + warn!( + ?worker_id, + ?operation_id, + "Action killed for memory with no reservation and no sampled peak; not escalating" + ); + return (update, false); + } + // With a ladder the next class up; without one, scale. Past the + // top class or the ceiling, the whole worker. + let next = if policy.ladder_kb.is_empty() { + Some( + base_kb + .saturating_mul(policy.percent) + .checked_div(100) + .unwrap_or(base_kb) + .max(base_kb.saturating_add(1)), + ) + } else { + policy + .ladder_kb + .iter() + .copied() + .filter(|kb| *kb > base_kb) + .min() + }; + let value = next.map_or(ceiling_kb, |kb| kb.min(ceiling_kb)); + let cpu = if value >= ceiling_kb { + whole_worker_cpu.map(|cores| (policy.cpu_property.clone(), cores)) + } else { + None + }; + let reason = result.error.clone().unwrap_or_else(|| { + make_err!( + Code::FailedPrecondition, + "Action was killed for memory at {} KiB", + usage.peak_memory_kb + ) + }); + warn!( + ?worker_id, + ?operation_id, + reserved_kb, + peak_memory_kb = usage.peak_memory_kb, + escalated_kb = value, + ceiling_kb, + property = %policy.property, + "Action killed for memory; requeuing with a larger reservation" + ); + ( + UpdateOperationType::UpdateWithEscalation(Escalation { + property: policy.property.clone(), + value, + reason, + cpu, + }), + true, + ) + } + async fn update_action( &mut self, worker_id: &WorkerId, @@ -713,7 +1030,11 @@ impl ApiWorkerSchedulerImpl { UpdateOperationType::UpdateWithError(err) => { (true, err.code == Code::ResourceExhausted) } - UpdateOperationType::UpdateWithDisconnect => (true, false), + // A decline pauses like backpressure: the worker said it has no + // room, so it gets nothing more until it says otherwise. + UpdateOperationType::UpdateWithDecline(_) => (true, true), + UpdateOperationType::UpdateWithDisconnect + | UpdateOperationType::UpdateWithEscalation(_) => (true, false), UpdateOperationType::ExecutionComplete => { // The process has exited but the action is still resident on // the worker until its result arrives, so nothing is released @@ -779,7 +1100,7 @@ impl ApiWorkerSchedulerImpl { complete_action_res }; - self.worker_change_notify.notify_one(); + self.capacity_changed(); update_operation_res.merge(complete_action_res) } @@ -791,16 +1112,46 @@ impl ApiWorkerSchedulerImpl { worker_id: WorkerId, operation_id: OperationId, action_info: ActionInfoWithProps, + dispatched_at: WorkerTimestamp, ) -> Result<(), Error> { if let Some(worker) = self.workers.get_mut(&worker_id) { let notify_worker_result = worker .notify_update(WorkerUpdate::RunAction(Box::new(( - operation_id, + operation_id.clone(), action_info.clone(), + dispatched_at, )))) .await; if let Err(notify_worker_result) = notify_worker_result { + if notify_worker_result.code == Code::ResourceExhausted { + // The worker's queue is full: it has stopped reading, or + // it is far behind. Nothing was charged; the action goes + // back untried and the worker waits for its next + // keepalive before it is offered anything else. + warn!( + ?worker_id, + %operation_id, + "Worker channel full, requeuing dispatch and pausing the worker" + ); + record_dispatch_requeue("channel_full"); + if !worker.is_paused { + worker.is_paused = true; + record_worker_state("paused", true); + } + self.worker_change_notify.notify_one(); + return self + .worker_state_manager + .update_operation( + &operation_id, + &worker_id, + UpdateOperationType::UpdateWithDecline(Decline { + reason: "channel_full".to_string(), + }), + ) + .await + .err_tip(|| "requeuing a dispatch the worker's channel refused"); + } warn!( ?worker_id, ?action_info, @@ -825,10 +1176,24 @@ impl ApiWorkerSchedulerImpl { } else { WorkerDisconnectReason::Error }; - return Result::<(), _>::Err(err.clone()).merge( - self.immediate_evict_worker(&worker_id, err, is_disconnect, reason) - .await, - ); + let evicted = self + .immediate_evict_worker(&worker_id, err.clone(), is_disconnect, reason) + .await; + // The dispatch was never charged to the worker (the send + // comes first), so the eviction did not see this operation; + // send it back to the queue here, untried: nothing ran, so + // no attempt is counted against it. + let requeued = self + .worker_state_manager + .update_operation( + &operation_id, + &worker_id, + UpdateOperationType::UpdateWithDecline(Decline { + reason: "worker_disconnected".to_string(), + }), + ) + .await; + return Result::<(), _>::Err(err).merge(evicted).merge(requeued); } Ok(()) } else { @@ -852,6 +1217,83 @@ impl ApiWorkerSchedulerImpl { /// Tells the worker to kill an operation it is still running but the /// state manager no longer has executing on it. A worker that cannot be /// reached is evicted, the same as for a failed run request. + /// The worker acknowledged a dispatch. A late acknowledgement for an + /// operation no longer on the worker is nothing to act on. + fn dispatch_accepted(&mut self, worker_id: &WorkerId, operation_id: &OperationId) { + let Some(worker) = self.workers.get_mut(worker_id) else { + return; + }; + if worker.mark_accepted(operation_id) { + return; + } + // Messages from a worker arrive in order and the acknowledgement + // comes before the run, so an acknowledgement for an operation not + // on this worker means the sweep took the dispatch back already: the + // worker is about to run an action that now belongs elsewhere. Tell + // it to stop now, not at the next revoked-operation sweep. + info!( + ?worker_id, + %operation_id, + "Acknowledgement for a dispatch already taken back; telling the worker to stop it" + ); + if let Err(err) = worker.kill_unknown_operation(operation_id) { + warn!(?worker_id, %operation_id, ?err, "Could not tell the worker to stop the operation"); + } + } + + /// The worker declined a dispatch: the ledger is restored, the action + /// requeued untried, and the worker paused; for a decline for load the + /// pause holds until the worker reports that much free memory. + /// Returns whether there was a dispatch to take back. A decline for an + /// operation the worker no longer holds (it finished, or the sweep + /// already requeued it) is nothing to act on, like a late + /// acknowledgement, and must not pause the worker. + async fn dispatch_declined( + &mut self, + worker_id: &WorkerId, + operation_id: &OperationId, + reason: String, + needs_kb: Option, + ) -> Result { + let holds_it = self + .workers + .peek(worker_id) + .is_some_and(|worker| worker.running_action_infos.contains_key(operation_id)); + if !holds_it { + debug!(?worker_id, %operation_id, %reason, "Decline for an operation not on this worker"); + return Ok(false); + } + info!(?worker_id, %operation_id, %reason, ?needs_kb, "Worker declined dispatch"); + self.update_action( + worker_id, + operation_id, + UpdateOperationType::UpdateWithDecline(Decline { reason }), + ) + .await?; + if let Some(worker) = self.workers.get_mut(worker_id) { + if !worker.is_paused { + worker.is_paused = true; + record_worker_state("paused", true); + } + worker.pause_needs_kb = needs_kb; + } + Ok(true) + } + + /// Dispatches this worker has not acknowledged within the timeout, as + /// of its latest keepalive. + fn unacknowledged_dispatches( + &self, + worker_id: &WorkerId, + timestamp: WorkerTimestamp, + timeout_s: u64, + ) -> Vec { + self.workers + .peek(worker_id) + .map(|worker| worker.unacknowledged(timestamp.saturating_sub(timeout_s))) + .unwrap_or_default() + } + async fn worker_notify_kill_operation( &mut self, worker_id: &WorkerId, @@ -876,6 +1318,18 @@ impl ApiWorkerSchedulerImpl { .notify_update(WorkerUpdate::KillOperation(operation_id.clone())) .await { + // A full channel is a worker that is behind, not one that is + // gone. It keeps its place; the request is forgotten so the + // next sweep sends the kill again. + if err.code == Code::ResourceExhausted { + worker.clear_kill_request(&operation_id); + warn!( + ?worker_id, + %operation_id, + "Worker channel full, kill deferred to the next sweep" + ); + return Ok(()); + } warn!( ?worker_id, %operation_id, @@ -924,11 +1378,11 @@ impl ApiWorkerSchedulerImpl { record_worker_disconnected(reason, worker.is_draining, worker.is_paused); // We don't care if we fail to send message to worker, this is only a best attempt. drop(worker.notify_update(WorkerUpdate::Disconnect).await); - let update = if is_disconnect { - UpdateOperationType::UpdateWithDisconnect - } else { - UpdateOperationType::UpdateWithError(err) - }; + // Whatever took the worker away (its stream ended, it timed + // out, a command to it failed, an operator removed it), the + // action running on it did nothing wrong: it is a worker loss + // with its own budget, not an attempt the action spent. + let update = UpdateOperationType::UpdateWithDisconnect; for (operation_id, _) in worker.running_action_infos.drain() { result = result.merge( self.worker_state_manager @@ -939,7 +1393,7 @@ impl ApiWorkerSchedulerImpl { } // Note: Calling this many time is very cheap, it'll only trigger `do_try_match` once. // TODO(palfrey) This should be moved to inside the Workers struct. - self.worker_change_notify.notify_one(); + self.capacity_changed(); result } } @@ -948,6 +1402,9 @@ impl ApiWorkerSchedulerImpl { pub struct ApiWorkerScheduler { #[metric] inner: Mutex, + /// See `ApiWorkerSchedulerImpl::capacity_generation`; read here without + /// the lock. + capacity_generation: Arc, #[metric(group = "platform_property_manager")] platform_property_manager: Arc, @@ -959,6 +1416,10 @@ pub struct ApiWorkerScheduler { help = "How long a sent kill may go unacknowledged before the worker is evicted, in seconds." )] unacknowledged_kill_timeout_s: u64, + #[metric( + help = "How long a dispatch may go unacknowledged before it is requeued untried, in seconds; 0 is off." + )] + dispatch_ack_timeout_s: u64, /// Shared worker registry for checking worker liveness. worker_registry: SharedWorkerRegistry, @@ -980,9 +1441,11 @@ impl ApiWorkerScheduler { platform_property_manager: Arc, allocation_strategy: WorkerAllocationStrategy, live_memory_veto: Option, + memory_escalation: Option, worker_change_notify: Arc, worker_timeout_s: u64, unacknowledged_kill_timeout_s: u64, + dispatch_ack_timeout_s: u64, worker_registry: SharedWorkerRegistry, maybe_origin_event_tx: Option>, has_peers: bool, @@ -993,13 +1456,16 @@ impl ApiWorkerScheduler { "record_ttl must be non-zero; a zero TTL degenerates the census warm-up/staleness checks" ); let local_fleet_change_notify = Arc::new(Notify::new()); + let capacity_generation = Arc::new(AtomicU64::new(0)); Arc::new(Self { inner: Mutex::new(ApiWorkerSchedulerImpl { workers: Workers(LruCache::unbounded()), worker_state_manager, allocation_strategy, live_memory_veto, + memory_escalation, worker_change_notify, + capacity_generation: capacity_generation.clone(), worker_registry: worker_registry.clone(), shutting_down: false, capability_index: WorkerCapabilityIndex::new(), @@ -1017,10 +1483,12 @@ impl ApiWorkerScheduler { platform_property_manager, worker_timeout_s, unacknowledged_kill_timeout_s, + dispatch_ack_timeout_s, worker_registry, metrics: Arc::new(SchedulerMetrics::default()), maybe_origin_event_tx, local_fleet_change_notify, + capacity_generation, }) } @@ -1034,26 +1502,92 @@ impl ApiWorkerScheduler { && inner.peer_fleet.as_ref().is_some_and(Vec::is_empty) } + /// The capacity generation as of now; see `find_worker_for_action_observed`. + pub fn capacity_generation(&self) -> u64 { + self.capacity_generation.load(Ordering::Relaxed) + } + /// Returns a reference to the worker registry. pub const fn worker_registry(&self) -> &SharedWorkerRegistry { &self.worker_registry } + /// `dispatched_at` is the scheduler's clock at the send; the + /// acknowledgement timeout counts from it. pub async fn worker_notify_run_action( &self, worker_id: WorkerId, operation_id: OperationId, action_info: ActionInfoWithProps, + dispatched_at: WorkerTimestamp, ) -> Result<(), Error> { self.metrics .actions_dispatched .fetch_add(1, Ordering::Relaxed); let mut inner = self.inner.lock().await; inner - .worker_notify_run_action(worker_id, operation_id, action_info) + .worker_notify_run_action(worker_id, operation_id, action_info, dispatched_at) .await } + /// A keepalive, or another message standing in for one. Only a + /// keepalive carries a load and lifts a pause; anything else refreshes + /// the timestamp and runs the acknowledgement sweep. + async fn refresh_worker( + &self, + worker_id: &WorkerId, + timestamp: WorkerTimestamp, + load: Option, + keepalive: bool, + ) -> Result<(), Error> { + { + let mut inner = self.inner.lock().await; + inner + .refresh_lifetime(worker_id, timestamp, load, keepalive) + .err_tip(|| "Error refreshing lifetime in worker_keep_alive_received()")?; + if self.dispatch_ack_timeout_s > 0 { + let stale = inner.unacknowledged_dispatches( + worker_id, + timestamp, + self.dispatch_ack_timeout_s, + ); + for operation_id in stale { + warn!( + ?worker_id, + %operation_id, + timeout_s = self.dispatch_ack_timeout_s, + "Dispatch never acknowledged, requeuing untried" + ); + match inner + .dispatch_declined( + worker_id, + &operation_id, + "unacknowledged".to_string(), + None, + ) + .await + { + Ok(true) => { + self.metrics + .dispatches_unacknowledged + .fetch_add(1, Ordering::Relaxed); + record_dispatch_requeue("unacknowledged"); + } + Ok(false) => {} + Err(err) => { + warn!(?worker_id, %operation_id, ?err, "Could not requeue unacknowledged dispatch"); + } + } + } + } + } + let now = UNIX_EPOCH + Duration::from_secs(timestamp); + self.worker_registry + .update_worker_heartbeat(worker_id, now) + .await; + Ok(()) + } + pub async fn running_action_info( &self, worker_id: &WorkerId, @@ -1109,6 +1643,22 @@ impl ApiWorkerScheduler { full_worker_logging: bool, now: SystemTime, ) -> MatchOutcome { + self.find_worker_for_action_observed(platform_properties, full_worker_logging, now) + .await + .0 + } + + /// `find_worker_for_action`, also returning the capacity generation read + /// under the same lock as the placement. A pass compares it with the one + /// it saw last: a change means a worker gained room since, and an older + /// action that found none earlier in the pass should be offered that room + /// before this one takes it. + pub async fn find_worker_for_action_observed( + &self, + platform_properties: &PlatformProperties, + full_worker_logging: bool, + now: SystemTime, + ) -> (MatchOutcome, u64) { let start = Instant::now(); self.metrics .find_worker_calls @@ -1118,6 +1668,7 @@ impl ApiWorkerScheduler { let worker_count = inner.workers.len() as u64; let result = inner.inner_find_worker_for_action(platform_properties, full_worker_logging, now); + let generation = inner.capacity_generation.load(Ordering::Relaxed); // Track workers iterated (worst case is all workers) self.metrics @@ -1138,7 +1689,7 @@ impl ApiWorkerScheduler { self.metrics .find_worker_time_ns .fetch_add(start.elapsed().as_nanos() as u64, Ordering::Relaxed); - result + (result, generation) } /// Checks to see if the worker exists in the worker pool. Should only be used in unit tests. @@ -1181,6 +1732,18 @@ impl WorkerScheduler for ApiWorkerScheduler { // default and `observed_worker_peak_memory_mib` was never recorded. // Sampling is optional, so only record when the worker actually took a // reading. A zero here means "not sampled", not "used no memory". + { + // The result follows this report; what it says decides whether + // that result is a kill to escalate. + let mut inner = self.inner.lock().await; + if let Some(pending) = inner + .workers + .get_mut(worker_id) + .and_then(|worker| worker.running_action_infos.get_mut(operation_id)) + { + pending.last_usage = Some(resource_usage.clone()); + } + } let maybe_action_info = self.running_action_info(worker_id, operation_id).await; let action_mnemonic = maybe_action_info .as_ref() @@ -1189,12 +1752,41 @@ impl WorkerScheduler for ApiWorkerScheduler { bazel_metadata.action_mnemonic.clone() }); + let outcome = + ResourceOutcome::try_from(resource_usage.outcome).unwrap_or(ResourceOutcome::Unknown); + if !matches!( + outcome, + ResourceOutcome::Unknown | ResourceOutcome::Completed + ) { + warn!( + ?worker_id, + ?operation_id, + outcome = outcome.as_str_name(), + enforced = resource_usage.enforced, + action_mnemonic, + peak_memory_kb = resource_usage.peak_memory_kb, + wall_time_ms = resource_usage.wall_time_ms, + reserved = ?resource_usage.reserved, + "Action ended by a kill" + ); + } + if resource_usage.sampled { if resource_usage.peak_memory_kb > 0 { - record_execution_peak_memory(resource_usage.peak_memory_kb, "", &action_mnemonic); + record_execution_peak_memory( + resource_usage.peak_memory_kb, + "", + &action_mnemonic, + outcome.as_str_name(), + ); } if resource_usage.cpu_time_ms > 0 { - record_execution_cpu_time(resource_usage.cpu_time_ms, "", &action_mnemonic); + record_execution_cpu_time( + resource_usage.cpu_time_ms, + "", + &action_mnemonic, + outcome.as_str_name(), + ); } } @@ -1275,28 +1867,61 @@ impl WorkerScheduler for ApiWorkerScheduler { update: UpdateOperationType, ) -> Result<(), Error> { let mut inner = self.inner.lock().await; + let (update, escalated) = inner.maybe_escalate(worker_id, operation_id, update); + if escalated { + record_dispatch_requeue("memory_escalation"); + } inner.update_action(worker_id, operation_id, update).await } - async fn worker_keep_alive_received( + async fn worker_dispatch_accepted( &self, worker_id: &WorkerId, - timestamp: WorkerTimestamp, - load: Option, + operation_id: &OperationId, ) -> Result<(), Error> { + let mut inner = self.inner.lock().await; + inner.dispatch_accepted(worker_id, operation_id); + Ok(()) + } + + async fn worker_dispatch_declined( + &self, + worker_id: &WorkerId, + operation_id: &OperationId, + reason: String, + needs_kb: Option, + ) -> Result<(), Error> { + let mut inner = self.inner.lock().await; + // Counted once the action is back in the queue, not per message. + if inner + .dispatch_declined(worker_id, operation_id, reason, needs_kb) + .await? { - let mut inner = self.inner.lock().await; - inner - .refresh_lifetime(worker_id, timestamp, load) - .err_tip(|| "Error refreshing lifetime in worker_keep_alive_received()")?; + self.metrics + .dispatches_declined + .fetch_add(1, Ordering::Relaxed); + record_dispatch_requeue("declined"); } - let now = UNIX_EPOCH + Duration::from_secs(timestamp); - self.worker_registry - .update_worker_heartbeat(worker_id, now) - .await; Ok(()) } + async fn worker_keep_alive_received( + &self, + worker_id: &WorkerId, + timestamp: WorkerTimestamp, + load: Option, + ) -> Result<(), Error> { + self.refresh_worker(worker_id, timestamp, load, true).await + } + + async fn worker_liveness_refreshed( + &self, + worker_id: &WorkerId, + timestamp: WorkerTimestamp, + ) -> Result<(), Error> { + self.refresh_worker(worker_id, timestamp, None, false).await + } + async fn remove_worker(&self, worker_id: &WorkerId) -> Result<(), Error> { self.worker_registry.remove_worker(worker_id).await; @@ -1446,6 +2071,42 @@ impl WorkerScheduler for ApiWorkerScheduler { result } + async fn worker_snapshot(&self) -> Vec { + fn as_strings(properties: &PlatformProperties) -> HashMap { + properties + .properties + .iter() + .map(|(name, value)| { + let value = match value { + PlatformPropertyValue::Minimum(v) => v.to_string(), + PlatformPropertyValue::Exact(v) + | PlatformPropertyValue::Priority(v) + | PlatformPropertyValue::Ignore(v) + | PlatformPropertyValue::Unknown(v) => v.clone(), + }; + (name.clone(), value) + }) + .collect() + } + let inner = self.inner.lock().await; + inner + .workers + .iter() + .map(|(_, worker)| WorkerSummary { + id: worker.id.to_string(), + running_actions: u32::try_from(worker.running_action_infos.len()) + .unwrap_or(u32::MAX), + max_inflight_tasks: worker.max_inflight_tasks, + is_paused: worker.is_paused, + is_draining: worker.is_draining, + last_update_timestamp: worker.last_update_timestamp, + platform_properties: as_strings(&worker.total_platform_properties), + available_platform_properties: as_strings(&worker.platform_properties), + free_memory_kb: worker.last_load.map(|load| load.free_memory_kb), + }) + .collect() + } + async fn set_drain_worker(&self, worker_id: &WorkerId, is_draining: bool) -> Result<(), Error> { let mut inner = self.inner.lock().await; inner.set_drain_worker(worker_id, is_draining) diff --git a/nativelink-scheduler/src/awaited_action_db/awaited_action.rs b/nativelink-scheduler/src/awaited_action_db/awaited_action.rs index 9c073338031..3adffbc1919 100644 --- a/nativelink-scheduler/src/awaited_action_db/awaited_action.rs +++ b/nativelink-scheduler/src/awaited_action_db/awaited_action.rs @@ -93,6 +93,19 @@ pub struct AwaitedAction { /// Number of attempts the job has been tried. #[metric(help = "The number of attempts the AwaitedAction has been tried")] pub attempts: usize, + + /// Requeues after the worker running the action was lost. Not an + /// attempt: the action did nothing wrong. Absent in records written + /// before the field existed. + #[serde(default)] + #[metric(help = "How many times the AwaitedAction lost its worker")] + pub worker_losses: usize, + + /// Requeues with a larger memory reservation after a memory kill. Not + /// an attempt either. Absent in older records. + #[serde(default)] + #[metric(help = "How many times the AwaitedAction was escalated")] + pub escalations: usize, } impl AwaitedAction { @@ -134,6 +147,8 @@ impl AwaitedAction { operation_id, sort_key, attempts: 0, + worker_losses: 0, + escalations: 0, last_worker_updated_timestamp: now, last_client_keepalive_timestamp: now, maybe_origin_metadata, @@ -210,6 +225,12 @@ impl AwaitedAction { self.last_client_keepalive_timestamp = now; } + /// Replaces the action's info, for an escalated reservation. The sort + /// key and identity stay: the action is the same one, asking for more. + pub(crate) fn set_action_info(&mut self, action_info: Arc) { + self.action_info = action_info; + } + pub(crate) fn set_client_operation_id(&mut self, client_operation_id: OperationId) { Arc::make_mut(&mut self.state).client_operation_id = client_operation_id; } diff --git a/nativelink-scheduler/src/default_scheduler_factory.rs b/nativelink-scheduler/src/default_scheduler_factory.rs index fef056c33f8..d30924eef86 100644 --- a/nativelink-scheduler/src/default_scheduler_factory.rs +++ b/nativelink-scheduler/src/default_scheduler_factory.rs @@ -16,7 +16,7 @@ use std::sync::Arc; use std::time::SystemTime; use nativelink_config::schedulers::{ - ExperimentalSimpleSchedulerBackend, SchedulerSpec, SimpleSpec, + ExperimentalSimpleSchedulerBackend, PropertyType, SchedulerSpec, SimpleSpec, }; use nativelink_config::stores::EvictionPolicy; use nativelink_error::{Error, ResultExt, make_input_err}; @@ -97,6 +97,36 @@ async fn inner_scheduler_factory( (Some(property_modifier_scheduler), worker_scheduler) } SchedulerSpec::HistoricalResource(spec) => { + HistoricalResourceScheduler::validate(spec)?; + // A number the scheduler writes into a property the nested + // scheduler does not treat as a minimum is matched as an exact + // string, and no worker advertises that string. + if let SchedulerSpec::Simple(simple) = spec.scheduler.as_ref() + && let Some(declared) = &simple.supported_platform_properties + { + let cold = spec.cold_start.unwrap_or_default(); + let dimensions = [ + ( + &spec.cpu_property_name, + spec.classes.iter().any(|c| c.cpu_count > 0) || cold.cpu_count > 0, + ), + ( + &spec.memory_property_name, + spec.classes.iter().any(|c| c.memory_kb > 0) || cold.memory_kb > 0, + ), + ( + &spec.disk_property_name, + spec.classes.iter().any(|c| c.disk_kb > 0) || cold.disk_kb > 0, + ), + ]; + for (name, used) in dimensions { + if used && declared.get(name) != Some(&PropertyType::Minimum) { + return Err(make_input_err!( + "historical_resource reserves {name} but the nested scheduler does not declare it as a minimum property; add it to supported_platform_properties as minimum" + )); + } + } + } let (action_scheduler, worker_scheduler) = Box::pin(inner_scheduler_factory( &spec.scheduler, store_manager, @@ -121,6 +151,14 @@ async fn simple_scheduler_factory( now_fn: fn() -> SystemTime, maybe_origin_event_tx: Option<&mpsc::Sender>, ) -> Result { + if let Some(policy) = &spec.memory_escalation + && policy.percent <= 100 + { + return Err(make_input_err!( + "memory_escalation.percent must be above 100 to grow the reservation, got {}", + policy.percent + )); + } match spec .experimental_backend .as_ref() @@ -166,7 +204,11 @@ async fn simple_scheduler_factory( task_change_notify.clone(), now_fn, Default::default, - spec.retain_completed_for_s, + // Passed through as 0, the store wrote every completed + // record with no expiry (the Lua update treats 0 as + // forever), so the scheduler's Redis kept every action ever + // run and the search index grew without bound. + retain_completed_for_s(spec.retain_completed_for_s), // Same normalisation SimpleScheduler applies, so the // keepalive stops being kept exactly when the timeout that // reads it comes due. @@ -190,8 +232,19 @@ async fn simple_scheduler_factory( } } +/// How long a completed record is kept, whichever backend holds it: the +/// configured value, or the same default for both when it is unset. +#[must_use] +pub const fn retain_completed_for_s(configured: u32) -> u32 { + if configured == 0 { + DEFAULT_RETAIN_COMPLETED_FOR_S + } else { + configured + } +} + pub fn memory_awaited_action_db_factory( - mut retain_completed_for_s: u32, + configured_retain_completed_for_s: u32, task_change_notify: &Arc, now_fn: NowFn, ) -> MemoryAwaitedActionDb @@ -199,12 +252,9 @@ where I: InstantWrapper, NowFn: Fn() -> I + Clone + Send + Sync + 'static, { - if retain_completed_for_s == 0 { - retain_completed_for_s = DEFAULT_RETAIN_COMPLETED_FOR_S; - } MemoryAwaitedActionDb::new( &EvictionPolicy { - max_seconds: retain_completed_for_s, + max_seconds: retain_completed_for_s(configured_retain_completed_for_s), ..Default::default() }, task_change_notify.clone(), diff --git a/nativelink-scheduler/src/historical_resource_scheduler.rs b/nativelink-scheduler/src/historical_resource_scheduler.rs index ad8bd0575ac..cc2b25f40a5 100644 --- a/nativelink-scheduler/src/historical_resource_scheduler.rs +++ b/nativelink-scheduler/src/historical_resource_scheduler.rs @@ -19,12 +19,13 @@ use std::sync::Arc; use std::time::Instant; use async_trait::async_trait; -use nativelink_config::schedulers::HistoricalResourceSpec; -use nativelink_error::{Error, ResultExt}; +use nativelink_config::schedulers::{HistoricalResourceSpec, SizeClassSpec, TimeoutClassSpec}; +use nativelink_error::{Error, ResultExt, make_input_err}; use nativelink_metric::{ MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent, RootMetricsComponent, }; use nativelink_util::action_messages::{ActionInfo, OperationId}; +use nativelink_util::metrics::record_hint_resolution; use nativelink_util::operation_state_manager::{ ActionStateResult, ActionStateResultStream, ClientStateManager, OperationFilter, }; @@ -37,12 +38,22 @@ use tracing::{debug, warn}; use crate::known_platform_property_provider::KnownPlatformPropertyProvider; -#[derive(Debug, Clone, Hash, Eq, PartialEq)] +#[derive(Debug, Clone, Default, Hash, Eq, PartialEq)] struct HintKey { target_id: Option, action_mnemonic: Option, + /// The action's `Command` digest as `hash-size`: an identity every + /// client has, whether or not it sends `RequestMetadata`. + command_digest: Option, + /// The `Action` digest as `hash-size`: what origin events carry, so + /// what a producer fed by them can key on. Narrower than the command + /// digest (the inputs are part of it). + action_digest: Option, } +/// One hint. The v1 fields keep their meaning; the v2 fields are optional, +/// so a v1 file still loads. `class` names a point on the configured +/// ladder and stands in for any number the hint leaves out. #[derive(Debug, Clone, Deserialize)] struct HistoricalResourceHint { #[serde(default)] @@ -50,22 +61,68 @@ struct HistoricalResourceHint { #[serde(default)] action_mnemonic: Option, #[serde(default)] + command_digest: Option, + #[serde(default)] + action_digest: Option, + #[serde(default)] cpu_count: Option, #[serde(default)] memory_kb: Option, #[serde(default)] memory_mib: Option, + #[serde(default)] + disk_kb: Option, + #[serde(default)] + class: Option, + /// Producer bookkeeping, carried for records. + #[serde(default)] + #[expect(dead_code, reason = "accepted from v2 files; not yet read")] + samples: Option, + #[serde(default)] + #[expect(dead_code, reason = "accepted from v2 files; not yet read")] + last_seen_s: Option, +} + +/// `hash/size` and `hash-size` both name the same digest. +fn normalize_digest_key(digest: &str) -> String { + digest.trim().replace('/', "-") } impl HistoricalResourceHint { - fn key(&self) -> Option { - if self.target_id.is_none() && self.action_mnemonic.is_none() { - return None; + /// Every key the lookup ladder could reach this hint by, one per + /// family it names: the target (with the mnemonic when it has one), + /// the action digest, the command digest, and the mnemonic alone when + /// there is no target. A record naming a target and a digest is found + /// by either; a single composite key would be found by neither. + fn keys(&self) -> Vec { + let digest = + |value: Option<&str>| value.map(normalize_digest_key).filter(|d| !d.is_empty()); + let mut keys = Vec::new(); + if self.target_id.is_some() { + keys.push(HintKey { + target_id: self.target_id.clone(), + action_mnemonic: self.action_mnemonic.clone(), + ..HintKey::default() + }); + } else if self.action_mnemonic.is_some() { + keys.push(HintKey { + action_mnemonic: self.action_mnemonic.clone(), + ..HintKey::default() + }); } - Some(HintKey { - target_id: self.target_id.clone(), - action_mnemonic: self.action_mnemonic.clone(), - }) + if let Some(action_digest) = digest(self.action_digest.as_deref()) { + keys.push(HintKey { + action_digest: Some(action_digest), + ..HintKey::default() + }); + } + if let Some(command_digest) = digest(self.command_digest.as_deref()) { + keys.push(HintKey { + command_digest: Some(command_digest), + ..HintKey::default() + }); + } + keys } fn memory_kb(&self) -> Option { @@ -97,11 +154,36 @@ struct HintState { last_loaded: Option, } +/// The numbers an action leaves `add_action` with, before they become +/// minimum properties. Zero means the dimension is not set. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +struct Numbers { + cpu_count: u64, + memory_kb: u64, + disk_kb: u64, +} + +impl From<&SizeClassSpec> for Numbers { + fn from(class: &SizeClassSpec) -> Self { + Self { + cpu_count: class.cpu_count, + memory_kb: class.memory_kb, + disk_kb: class.disk_kb, + } + } +} + pub struct HistoricalResourceScheduler { hints_file: String, refresh_interval: Duration, cpu_property_name: String, memory_property_name: String, + disk_property_name: String, + cold_start: Numbers, + classes: Vec, + default_class: Option, + class_by_mnemonic: HashMap, + class_by_timeout: Vec, scheduler: Option>, known_properties: Mutex>>, hint_state: Mutex, @@ -114,6 +196,9 @@ impl core::fmt::Debug for HistoricalResourceScheduler { .field("refresh_interval", &self.refresh_interval) .field("cpu_property_name", &self.cpu_property_name) .field("memory_property_name", &self.memory_property_name) + .field("cold_start", &self.cold_start) + .field("classes", &self.classes) + .field("default_class", &self.default_class) .field("known_properties", &self.known_properties) .finish_non_exhaustive() } @@ -130,6 +215,19 @@ impl HistoricalResourceScheduler { refresh_interval: Duration::from_secs(spec.refresh_interval_s), cpu_property_name: spec.cpu_property_name.clone(), memory_property_name: spec.memory_property_name.clone(), + disk_property_name: spec.disk_property_name.clone(), + cold_start: spec + .cold_start + .as_ref() + .map_or_else(Numbers::default, |c| Numbers { + cpu_count: c.cpu_count, + memory_kb: c.memory_kb, + disk_kb: c.disk_kb, + }), + classes: spec.classes.clone(), + default_class: spec.default_class.clone(), + class_by_mnemonic: spec.class_by_mnemonic.clone(), + class_by_timeout: spec.class_by_timeout.clone(), scheduler: Some(scheduler), known_properties: Mutex::new(HashMap::new()), hint_state: Mutex::default(), @@ -164,9 +262,34 @@ impl HistoricalResourceScheduler { return; } }; + // A class the ladder does not have is dropped here, once per name + // per load, so the hot path never logs it; the hint keeps whatever + // numbers it states. + let mut unknown_classes = HashSet::new(); + let mut hints: Vec = hints + .into_iter() + .map(|mut hint| { + if let Some(name) = &hint.class + && self.class(name).is_none() + { + if unknown_classes.insert(name.clone()) { + warn!( + class = %name, + hints_file = %self.hints_file, + "Hints name a size class the ladder does not have; the class is ignored" + ); + } + hint.class = None; + } + hint + }) + .collect(); + // A record naming fewer families is the more general one, so it + // owns a shared key: sort so those are inserted last. + hints.sort_by_key(|hint| core::cmp::Reverse(hint.keys().len())); let hints = hints .into_iter() - .filter_map(|hint| hint.key().map(|key| (key, hint))) + .flat_map(|hint| hint.keys().into_iter().map(move |key| (key, hint.clone()))) .collect::>(); debug!(hints_file = %self.hints_file, hints = hints.len(), "Loaded historical resource hints"); let mut hint_state = self.hint_state.lock(); @@ -174,53 +297,190 @@ impl HistoricalResourceScheduler { hint_state.last_loaded = Some(now); } - fn hint_for_current_action(&self) -> Option { - self.refresh_hints(); + /// The Bazel request metadata for the current action, if the client + /// sent any: `(target_id, action_mnemonic)`, each empty string as `None`. + fn current_metadata() -> (Option, Option) { let ctx = Context::current(); - let metadata = ctx + let Some(metadata) = ctx .baggage() .get(BAZEL_METADATA_KEY) - .and_then(|value| request_metadata_from_baggage(value.as_str()).ok())?; - let target_id = non_empty_string(metadata.target_id); - let action_mnemonic = non_empty_string(metadata.action_mnemonic); - let hint_state = self.hint_state.lock(); - let exact_key = HintKey { - target_id: target_id.clone(), - action_mnemonic: action_mnemonic.clone(), + .and_then(|value| request_metadata_from_baggage(value.as_str()).ok()) + else { + return (None, None); }; - hint_state - .hints - .get(&exact_key) - .or_else(|| { - hint_state.hints.get(&HintKey { - target_id, - action_mnemonic: None, - }) - }) - .or_else(|| { - hint_state.hints.get(&HintKey { - target_id: None, - action_mnemonic, - }) + ( + non_empty_string(metadata.target_id), + non_empty_string(metadata.action_mnemonic), + ) + } + + /// The key ladder, most specific first: `(target, mnemonic)`, then + /// `target`, then the action digest, then the command digest, then the + /// mnemonic alone. The digest keys are what make hints reach clients + /// that send no `RequestMetadata`, and the same command regardless of + /// who runs it. + fn hint_for_action( + &self, + action_info: &ActionInfo, + target_id: Option<&String>, + action_mnemonic: Option<&String>, + ) -> Option<(HistoricalResourceHint, &'static str)> { + self.refresh_hints(); + let hint_state = self.hint_state.lock(); + if hint_state.hints.is_empty() { + return None; + } + let by = |target_id: Option<&String>, + action_mnemonic: Option<&String>, + command_digest: Option, + action_digest: Option| { + hint_state.hints.get(&HintKey { + target_id: target_id.cloned(), + action_mnemonic: action_mnemonic.cloned(), + command_digest, + action_digest, }) - .cloned() + }; + if target_id.is_some() + && action_mnemonic.is_some() + && let Some(hint) = by(target_id, action_mnemonic, None, None) + { + return Some((hint.clone(), "target")); + } + if target_id.is_some() + && let Some(hint) = by(target_id, None, None, None) + { + return Some((hint.clone(), "target")); + } + if let Some(hint) = by(None, None, None, Some(action_info.digest().to_string())) { + return Some((hint.clone(), "action_digest")); + } + if let Some(hint) = by( + None, + None, + Some(action_info.command_digest.to_string()), + None, + ) { + return Some((hint.clone(), "command_digest")); + } + if action_mnemonic.is_some() + && let Some(hint) = by(None, action_mnemonic, None, None) + { + return Some((hint.clone(), "mnemonic")); + } + None } - fn apply_hint(&self, action_info: &mut ActionInfo, hint: &HistoricalResourceHint) { - if let Some(cpu_count) = hint.cpu_count { - apply_minimum_platform_property( - &mut action_info.platform_properties, - &self.cpu_property_name, - cpu_count, - ); + fn class(&self, name: &str) -> Option<&SizeClassSpec> { + self.classes.iter().find(|class| class.name == name) + } + + /// A hint's numbers: what it states, with its class filling anything + /// it leaves out. An unknown class was dropped at load. + fn numbers_from_hint(&self, hint: &HistoricalResourceHint) -> Numbers { + let from_class = hint + .class + .as_deref() + .and_then(|name| self.class(name)) + .map(Numbers::from) + .unwrap_or_default(); + Numbers { + cpu_count: hint.cpu_count.unwrap_or(from_class.cpu_count), + memory_kb: hint.memory_kb().unwrap_or(from_class.memory_kb), + disk_kb: hint.disk_kb.unwrap_or(from_class.disk_kb), } - if let Some(memory_kb) = hint.memory_kb() { - apply_minimum_platform_property( - &mut action_info.platform_properties, - &self.memory_property_name, - memory_kb, + } + + /// Raises the action's minimums to `numbers`; never lowers a value the + /// client or an earlier step set. + fn apply_minimums(&self, action_info: &mut ActionInfo, numbers: Numbers) { + apply_minimum_platform_property( + &mut action_info.platform_properties, + &self.cpu_property_name, + numbers.cpu_count, + ); + apply_minimum_platform_property( + &mut action_info.platform_properties, + &self.memory_property_name, + numbers.memory_kb, + ); + apply_minimum_platform_property( + &mut action_info.platform_properties, + &self.disk_property_name, + numbers.disk_kb, + ); + } + + /// The cold-start policy for an action no hint matched: a class by + /// mnemonic, else by timeout, else the default class; without classes, + /// the raw `cold_start` numbers. Only absent properties are filled. + fn cold_start_numbers( + &self, + action_info: &ActionInfo, + action_mnemonic: Option<&String>, + ) -> Numbers { + let by_mnemonic = action_mnemonic.and_then(|m| self.class_by_mnemonic.get(m)); + let by_timeout = if action_info.timeout.is_zero() { + None + } else { + self.class_by_timeout + .iter() + .filter(|rule| Duration::from_secs(rule.min_timeout_s) <= action_info.timeout) + .max_by_key(|rule| rule.min_timeout_s) + .map(|rule| &rule.class) + }; + // The rules were checked against the ladder at load; the first + // candidate that names a class on it wins. + [by_mnemonic, by_timeout, self.default_class.as_ref()] + .into_iter() + .flatten() + .find_map(|name| self.class(name)) + .map_or(self.cold_start, Numbers::from) + } + + /// The cold-start rules and the default class must name classes on the + /// ladder, or the action they were written for gets nothing. + pub fn validate(spec: &HistoricalResourceSpec) -> Result<(), Error> { + let known = |name: &str| spec.classes.iter().any(|class| class.name == name); + let mut rules = spec + .class_by_mnemonic + .iter() + .map(|(mnemonic, class)| (format!("class_by_mnemonic[{mnemonic}]"), class.as_str())) + .chain(spec.class_by_timeout.iter().map(|rule| { + ( + format!("class_by_timeout[{}]", rule.min_timeout_s), + rule.class.as_str(), + ) + })) + .chain( + spec.default_class + .iter() + .map(|class| ("default_class".to_string(), class.as_str())), ); + if let Some((rule, class)) = rules.find(|(_, class)| !known(class)) { + return Err(make_input_err!( + "historical_resource.{rule} names size class {class:?}, which classes does not define" + )); } + Ok(()) + } + + fn apply_defaults(&self, action_info: &mut ActionInfo, numbers: Numbers) { + apply_default_platform_property( + &mut action_info.platform_properties, + &self.cpu_property_name, + numbers.cpu_count, + ); + apply_default_platform_property( + &mut action_info.platform_properties, + &self.memory_property_name, + numbers.memory_kb, + ); + apply_default_platform_property( + &mut action_info.platform_properties, + &self.disk_property_name, + numbers.disk_kb, + ); } async fn inner_get_known_properties(&self, instance_name: &str) -> Result, Error> { @@ -241,6 +501,7 @@ impl HistoricalResourceScheduler { ); known_properties.insert(self.cpu_property_name.clone()); known_properties.insert(self.memory_property_name.clone()); + known_properties.insert(self.disk_property_name.clone()); let final_known_properties: Vec = known_properties.into_iter().collect(); self.known_properties .lock() @@ -270,6 +531,17 @@ fn apply_minimum_platform_property( } } +fn apply_default_platform_property( + platform_properties: &mut HashMap, + property_name: &str, + default_value: u64, +) { + if default_value == 0 || platform_properties.contains_key(property_name) { + return; + } + platform_properties.insert(property_name.to_string(), default_value.to_string()); +} + #[async_trait] impl KnownPlatformPropertyProvider for HistoricalResourceScheduler { async fn get_known_properties(&self, instance_name: &str) -> Result, Error> { @@ -284,8 +556,25 @@ impl ClientStateManager for HistoricalResourceScheduler { client_operation_id: OperationId, mut action_info: Arc, ) -> Result, Error> { - if let Some(hint) = self.hint_for_current_action() { - self.apply_hint(Arc::make_mut(&mut action_info), &hint); + let (target_id, action_mnemonic) = Self::current_metadata(); + // A hint that comes out all zero (bookkeeping fields only, or a + // class that was dropped at load) reserves nothing, so it is no + // hint: the cold-start policy applies and the metric says so. + let hint = self + .hint_for_action(&action_info, target_id.as_ref(), action_mnemonic.as_ref()) + .map(|(hint, source)| (self.numbers_from_hint(&hint), source)) + .filter(|(numbers, _)| *numbers != Numbers::default()); + if let Some((numbers, source)) = hint { + self.apply_minimums(Arc::make_mut(&mut action_info), numbers); + record_hint_resolution(source); + } else { + let numbers = self.cold_start_numbers(&action_info, action_mnemonic.as_ref()); + if numbers == Numbers::default() { + record_hint_resolution("none"); + } else { + self.apply_defaults(Arc::make_mut(&mut action_info), numbers); + record_hint_resolution("cold_start"); + } } self.scheduler .as_ref() diff --git a/nativelink-scheduler/src/lib.rs b/nativelink-scheduler/src/lib.rs index dfdc2f53c9d..40fe60d0feb 100644 --- a/nativelink-scheduler/src/lib.rs +++ b/nativelink-scheduler/src/lib.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +pub mod admin; pub mod api_worker_scheduler; pub mod awaited_action_db; pub mod cache_lookup_scheduler; @@ -32,3 +33,5 @@ pub mod worker; pub mod worker_capability_index; pub mod worker_registry; pub mod worker_scheduler; + +// CI affected-lanes validation: scheduler-only touch. diff --git a/nativelink-scheduler/src/simple_scheduler.rs b/nativelink-scheduler/src/simple_scheduler.rs index 48842ad2d21..9a20ba2030a 100644 --- a/nativelink-scheduler/src/simple_scheduler.rs +++ b/nativelink-scheduler/src/simple_scheduler.rs @@ -13,9 +13,9 @@ // limitations under the License. use core::sync::atomic::{AtomicUsize, Ordering}; -use std::collections::{BTreeSet, HashMap}; +use std::collections::{BTreeSet, HashMap, HashSet, VecDeque}; use std::sync::Arc; -use std::time::{Instant, SystemTime}; +use std::time::{Instant, SystemTime, UNIX_EPOCH}; use async_trait::async_trait; use futures::{Future, StreamExt, future}; @@ -31,8 +31,9 @@ use nativelink_proto::com::github::trace_machina::nativelink::remote_execution:: use nativelink_util::action_messages::{ActionInfo, ActionState, OperationId, WorkerId}; use nativelink_util::instant_wrapper::InstantWrapper; use nativelink_util::metrics::{ - record_awaited_action_orphan, record_matching_pass, record_queue_depth, record_sweep_failure, - record_unsatisfiable_failed, record_unsatisfiable_queued, + record_awaited_action_orphan, record_matching_pass, record_parked_dispatched, + record_queue_depth, record_sweep_failure, record_unsatisfiable_failed, + record_unsatisfiable_queued, }; use nativelink_util::operation_state_manager::{ ActionStateResult, ActionStateResultStream, ClientStateManager, MatchingEngineStateManager, @@ -58,11 +59,13 @@ use crate::awaited_action_db::{AwaitedActionDb, CLIENT_KEEPALIVE_DURATION}; use crate::known_platform_property_provider::KnownPlatformPropertyProvider; use crate::match_outcome::{MatchOutcome, PropertyShape, UnsatisfiableReason}; use crate::platform_property_manager::PlatformPropertyManager; -use crate::simple_scheduler_state_manager::{SimpleSchedulerStateManager, is_lost_record}; +use crate::simple_scheduler_state_manager::{ + RetryLimits, SimpleSchedulerStateManager, is_lost_record, +}; use crate::unsatisfiable_tracker::UnsatisfiableTracker; use crate::worker::{ActionInfoWithProps, Worker, WorkerTimestamp}; use crate::worker_registry::WorkerRegistry; -use crate::worker_scheduler::WorkerScheduler; +use crate::worker_scheduler::{WorkerScheduler, WorkerSummary}; /// Default timeout for workers in seconds: four default keepalive intervals /// (2.5s each). The old value of two meant one late keepalive under load @@ -214,6 +217,14 @@ pub struct SimpleScheduler { /// (`no_worker_action_timeout_s`); `None` when that timeout is off. no_worker_tracker: Option>, + /// Typed platform properties by the hash of the raw ones. Typing is a + /// lookup and a parse per property per queued action per pass, and the + /// raw properties of a queued action rarely change (an escalation is + /// the exception, and it changes the hash), so the pass reuses the + /// typed set. Cleared when full rather than evicted, like + /// `static_verdicts`. + typed_properties: Mutex>, + /// The number of unsatisfiable actions still queued at the end of the last /// matching pass, used to size the next pass's per-pass fail cap so a large /// burst drains in a few passes. See `unsatisfiable_fail_cap`. @@ -242,6 +253,37 @@ impl core::fmt::Debug for SimpleScheduler { } } +/// How many typed property sets the pass keeps before starting over. +const MAX_TYPED_PROPERTIES: usize = 10_000; + +/// The typed form of an action's raw properties, from the cache when the +/// same raw set was typed before. +fn typed_properties_for( + cache: &Mutex>, + manager: &PlatformPropertyManager, + raw: &HashMap, +) -> Result { + use core::hash::{Hash, Hasher}; + let mut pairs: Vec<(&String, &String)> = raw.iter().collect(); + pairs.sort(); + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + for (name, value) in &pairs { + name.hash(&mut hasher); + value.hash(&mut hasher); + } + let key = hasher.finish(); + if let Some(typed) = cache.lock().get(&key) { + return Ok(typed.clone()); + } + let typed = manager.make_platform_properties(raw.clone())?; + let mut cache = cache.lock(); + if cache.len() >= MAX_TYPED_PROPERTIES { + cache.clear(); + } + cache.insert(key, typed.clone()); + Ok(typed) +} + impl SimpleScheduler { fn origin_event_id(event: &Event) -> String { Uuid::now_v6(&get_node_id(Some(event))) @@ -514,27 +556,26 @@ impl SimpleScheduler { Ok(()) } - #[expect(clippy::too_many_arguments)] - async fn match_action_to_worker( - action_state_result: &dyn ActionStateResult, - workers: &ApiWorkerScheduler, - matching_engine_state_manager: &dyn MatchingEngineStateManager, + /// A listed action with what matching needs of it. + struct Loaded { + action_state_result: Box, + action_info: ActionInfoWithProps, + } + + /// Reads what matching needs of a listed action. `None` when the + /// record cannot be read: gone from the store (eviction), or + /// unreadable. Counted so an operator can see the store losing + /// records, skipped so the rest of the pass still runs, and not an + /// error, since an error here made the pass rerun at once and log + /// ten times a second until the entry went away. + async fn load_action( + action_state_result: Box, platform_property_manager: &PlatformPropertyManager, - maybe_origin_event_tx: Option<&mpsc::Sender>, - full_worker_logging: bool, - unsatisfiable_pass: &UnsatisfiablePass<'_>, - no_worker_pass: Option<&UnsatisfiablePass<'_>>, - ) -> Result<(), Error> { + typed_properties: &Mutex>, + ) -> Result, Error> { let (action_info, maybe_origin_metadata) = match action_state_result.as_action_info().await { Ok(found) => found, - // Listed by the queue, and its record is gone from the - // store (eviction) or cannot be decoded. Nothing to - // match. Counted so an operator can see the store losing - // records, skipped so the rest of the pass still runs, - // and not an error, since an error here made the pass - // rerun at once and log ten times a second until the - // entry went away. Err(err) if is_lost_record(&err) => { if err.code == Code::NotFound { debug!( @@ -548,7 +589,7 @@ impl SimpleScheduler { ); } record_awaited_action_orphan("matching"); - return Ok(()); + return Ok(None); } // The store itself failed, so the pass is incomplete: // an error here reruns it at once, as before. @@ -559,70 +600,46 @@ impl SimpleScheduler { } }; - // TODO(palfrey) We should not compute this every time and instead store - // it with the ActionInfo when we receive it. - let platform_properties = platform_property_manager - .make_platform_properties(action_info.platform_properties.clone()) - .err_tip( - || "Failed to make platform properties in SimpleScheduler::do_try_match", - )?; - - let origin_metadata = maybe_origin_metadata.unwrap_or_default(); - let action_info = ActionInfoWithProps { - inner: action_info, - platform_properties, - origin_metadata: origin_metadata.clone(), - scheduler_start_execute_event_id: None, - }; + let platform_properties = typed_properties_for( + typed_properties, + platform_property_manager, + &action_info.platform_properties, + ) + .err_tip(|| "Failed to make platform properties in SimpleScheduler::do_try_match")?; + + Ok(Some(Loaded { + action_state_result, + action_info: ActionInfoWithProps { + inner: action_info, + platform_properties, + origin_metadata: maybe_origin_metadata.unwrap_or_default(), + scheduler_start_execute_event_id: None, + }, + })) + } - // Try to find a worker for the action. - let worker_id = { - match workers - .find_worker_for_action( - &action_info.platform_properties, - full_worker_logging, - unsatisfiable_pass.now, - ) - .await - { - MatchOutcome::Matched(worker_id) => worker_id, - // The action can run once a worker has room, so we have - // nothing to do. - MatchOutcome::WaitingForCapacity => { - return Ok(()); - } - // Nothing to wait for; a worker connecting starts a new - // pass, unless the no-worker clock is running. - MatchOutcome::NoWorkersConnected => { - let Some(no_worker_pass) = no_worker_pass else { - return Ok(()); - }; - return handle_unsatisfiable( - action_state_result, - matching_engine_state_manager, - workers, - &action_info.platform_properties, - action_info.inner.insert_timestamp, - &UnsatisfiableReason::no_workers(), - no_worker_pass, - ) - .await; - } - MatchOutcome::Unsatisfiable(reason) => { - return handle_unsatisfiable( - action_state_result, - matching_engine_state_manager, - workers, - &action_info.platform_properties, - action_info.inner.insert_timestamp, - &reason, - unsatisfiable_pass, - ) - .await; - } - } - }; + /// What became of a dispatch that did not fail. + #[derive(Clone, Copy, PartialEq, Eq)] + enum Dispatch { + /// The worker was told to run the action. + Sent, + /// Another scheduler assigned the action first; nothing was sent. + LostRace, + } + /// Hands a matched action to its worker. `Ok(Dispatch::Sent)` means + /// the worker was actually told; a lost assignment race is not an + /// error but sends nothing. + async fn dispatch_to_worker( + action_state_result: &dyn ActionStateResult, + workers: &ApiWorkerScheduler, + matching_engine_state_manager: &dyn MatchingEngineStateManager, + maybe_origin_event_tx: Option<&mpsc::Sender>, + worker_id: WorkerId, + action_info: ActionInfoWithProps, + now_fn: &(dyn Fn() -> SystemTime + Send + Sync), + ) -> Result { + let origin_metadata = action_info.origin_metadata.clone(); let event_origin_metadata = origin_metadata.clone(); let attach_operation_fut = async move { // Extract the operation_id from the action_state. @@ -643,7 +660,7 @@ impl SimpleScheduler { if err.code == Code::Aborted { // If the operation was aborted, it means that the operation was // cancelled due to another operation being assigned to the worker. - return Ok(()); + return Ok(Dispatch::LostRace); } // Any other error is a real error. return Err(err); @@ -662,8 +679,12 @@ impl SimpleScheduler { }); debug!(%worker_id, %operation_id, ?action_info, "Notifying worker of operation"); + let dispatched_at = now_fn() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); workers - .worker_notify_run_action(worker_id, operation_id, action_info) + .worker_notify_run_action(worker_id, operation_id, action_info, dispatched_at) .await .err_tip(|| { "Failed to run worker_notify_run_action in SimpleScheduler::do_try_match" @@ -679,7 +700,7 @@ impl SimpleScheduler { .await; } - Ok(()) + Ok(Dispatch::Sent) }; tokio::pin!(attach_operation_fut); @@ -694,6 +715,114 @@ impl SimpleScheduler { .await } + /// Actions this pass read that found no room, in the order the + /// queue listed them: highest priority first, then oldest. That is + /// the order room is offered to them in, before the action the + /// pass is reading, whenever the fleet's capacity changes. Parking + /// is a push; a replay walks the list once and asks the fleet once + /// per property shape it meets, since every action of one shape + /// asks the same of a worker: a shape whose first parked action + /// finds no room is skipped for the rest of that replay. + #[derive(Default)] + struct Parked { + /// Each with its shape, computed once at parking. + queue: VecDeque<(PropertyShape, Loaded)>, + /// Actions that were not parked because the list was full. + cap_dropped: u64, + } + + /// Past this many parked actions a pass stops parking, so a newer + /// action past that point may take room ahead of one of them. The + /// head of the queue, which is what a client waits on, is always + /// parked. + const PARKED_CAP: usize = 10_000; + + impl Parked { + fn park(&mut self, loaded: Loaded) { + if self.queue.len() >= PARKED_CAP { + self.cap_dropped += 1; + return; + } + let shape = PropertyShape::from(&loaded.action_info.platform_properties); + self.queue.push_back((shape, loaded)); + } + + fn is_empty(&self) -> bool { + self.queue.is_empty() + } + } + + /// What one pass did with its parked actions. + #[derive(Default)] + struct Placement { + /// Parked actions that found a worker and were dispatched to it. + attempted: u64, + /// Of those, the ones the worker was actually told to run. + dispatched: u64, + /// Of those, the ones another scheduler assigned first. + lost_race: u64, + /// Of those, the ones whose dispatch failed. + failed: u64, + } + + /// Offers the room that opened since the pass last looked to the + /// parked actions in listing order. Every parked action that found + /// a worker leaves the list, whatever its dispatch came to. + #[expect(clippy::too_many_arguments)] + async fn place_parked( + parked: &mut Parked, + placement: &mut Placement, + workers: &ApiWorkerScheduler, + matching_engine_state_manager: &dyn MatchingEngineStateManager, + maybe_origin_event_tx: Option<&mpsc::Sender>, + full_worker_logging: bool, + now: SystemTime, + now_fn: &(dyn Fn() -> SystemTime + Send + Sync), + ) -> Result<(), Error> { + let mut result = Ok(()); + let mut blocked: HashSet = HashSet::new(); + let mut still_parked = VecDeque::with_capacity(parked.queue.len()); + while let Some((shape, loaded)) = parked.queue.pop_front() { + if blocked.contains(&shape) { + still_parked.push_back((shape, loaded)); + continue; + } + let MatchOutcome::Matched(worker_id) = workers + .find_worker_for_action( + &loaded.action_info.platform_properties, + full_worker_logging, + now, + ) + .await + else { + still_parked.push_back((shape.clone(), loaded)); + blocked.insert(shape); + continue; + }; + placement.attempted += 1; + match dispatch_to_worker( + loaded.action_state_result.as_ref(), + workers, + matching_engine_state_manager, + maybe_origin_event_tx, + worker_id, + loaded.action_info, + now_fn, + ) + .await + { + Ok(Dispatch::Sent) => placement.dispatched += 1, + Ok(Dispatch::LostRace) => placement.lost_race += 1, + Err(err) => { + placement.failed += 1; + result = result.merge(Err(err)); + } + } + } + parked.queue = still_parked; + result + } + let mut result = Ok(()); let start = Instant::now(); @@ -725,6 +854,7 @@ impl SimpleScheduler { } _ => None, }; + let no_workers_reason = UnsatisfiableReason::no_workers(); let mut stream = self .get_queued_operations() @@ -739,23 +869,143 @@ impl SimpleScheduler { ); } + let workers = self.worker_scheduler.as_ref(); + let matching_engine_state_manager = self.matching_engine_state_manager.as_ref(); + let maybe_origin_event_tx = self.maybe_origin_event_tx.as_ref(); + let mut parked = Parked::default(); + let mut placement = Placement::default(); + let mut capacity_seen = workers.capacity_generation(); let mut queued_seen: u64 = 0; while let Some(action_state_result) = stream.next().await { queued_seen += 1; + let loaded = match load_action( + action_state_result, + self.platform_property_manager.as_ref(), + &self.typed_properties, + ) + .await + { + Ok(Some(loaded)) => loaded, + Ok(None) => continue, + Err(err) => { + result = result.merge(Err(err)); + continue; + } + }; + + // Room that opened since the pass last placed anything belongs + // to the oldest action still waiting, not to this one, which is + // newer than everything parked. The generation is read under the + // placement lock, so a change cannot slip between the two. + let outcome = loop { + let (outcome, generation) = workers + .find_worker_for_action_observed( + &loaded.action_info.platform_properties, + full_worker_logging, + unsatisfiable_pass.now, + ) + .await; + if generation == capacity_seen || parked.is_empty() { + capacity_seen = generation; + break outcome; + } + capacity_seen = generation; + result = result.merge( + place_parked( + &mut parked, + &mut placement, + workers, + matching_engine_state_manager, + maybe_origin_event_tx, + full_worker_logging, + unsatisfiable_pass.now, + self.now_fn.as_ref(), + ) + .await, + ); + }; + + match outcome { + MatchOutcome::Matched(worker_id) => { + result = result.merge( + dispatch_to_worker( + loaded.action_state_result.as_ref(), + workers, + matching_engine_state_manager, + maybe_origin_event_tx, + worker_id, + loaded.action_info, + self.now_fn.as_ref(), + ) + .await + .map(|_| ()), + ); + } + // A worker could run it once one has room. + MatchOutcome::WaitingForCapacity => parked.park(loaded), + // Nothing to wait for; a worker connecting starts a new pass, + // unless the no-worker clock is running. + MatchOutcome::NoWorkersConnected => { + if let Some(no_worker_pass) = &no_worker_pass { + result = result.merge( + handle_unsatisfiable( + loaded.action_state_result.as_ref(), + matching_engine_state_manager, + workers, + &loaded.action_info.platform_properties, + loaded.action_info.inner.insert_timestamp, + &no_workers_reason, + no_worker_pass, + ) + .await, + ); + } + } + MatchOutcome::Unsatisfiable(reason) => { + result = result.merge( + handle_unsatisfiable( + loaded.action_state_result.as_ref(), + matching_engine_state_manager, + workers, + &loaded.action_info.platform_properties, + loaded.action_info.inner.insert_timestamp, + &reason, + &unsatisfiable_pass, + ) + .await, + ); + } + } + } + // Room that opened after the last action was read. + if !parked.is_empty() && workers.capacity_generation() != capacity_seen { result = result.merge( - match_action_to_worker( - action_state_result.as_ref(), - self.worker_scheduler.as_ref(), - self.matching_engine_state_manager.as_ref(), - self.platform_property_manager.as_ref(), - self.maybe_origin_event_tx.as_ref(), + place_parked( + &mut parked, + &mut placement, + workers, + matching_engine_state_manager, + maybe_origin_event_tx, full_worker_logging, - &unsatisfiable_pass, - no_worker_pass.as_ref(), + unsatisfiable_pass.now, + self.now_fn.as_ref(), ) .await, ); } + record_parked_dispatched(placement.dispatched); + if placement.attempted > 0 || parked.cap_dropped > 0 { + debug!( + attempted = placement.attempted, + dispatched = placement.dispatched, + lost_race = placement.lost_race, + failed = placement.failed, + cap_dropped = parked.cap_dropped, + still_parked = parked.queue.len(), + elapsed_ms = start.elapsed().as_millis(), + "Parked actions offered the room that opened mid-pass" + ); + } let mut unsatisfiable_queued = self .unsatisfiable_tracker @@ -926,8 +1176,19 @@ impl SimpleScheduler { let record_ttl = fleet_exchange_interval * FLEET_RECORD_TTL_INTERVALS; let has_peers = awaited_action_db.shares_state(); - let state_manager = SimpleSchedulerStateManager::new( + let retry_limits = RetryLimits { max_job_retries, + max_worker_loss_retries: if spec.max_worker_loss_retries == 0 { + RetryLimits::DEFAULT_WORKER_LOSS_RETRIES + } else { + spec.max_worker_loss_retries + }, + max_escalations: spec.memory_escalation.as_ref().map_or(0, |policy| { + usize::try_from(policy.max_steps).unwrap_or(usize::MAX) + }), + }; + let state_manager = SimpleSchedulerStateManager::new( + retry_limits, Duration::from_secs(worker_timeout_s), Duration::from_secs(client_action_timeout_s), Duration::from_secs(spec.max_action_executing_timeout_s), @@ -941,9 +1202,11 @@ impl SimpleScheduler { platform_property_manager.clone(), spec.allocation_strategy, spec.live_memory_veto.clone(), + spec.memory_escalation.clone(), worker_change_notify.clone(), worker_timeout_s, unacknowledged_kill_timeout_s, + spec.dispatch_ack_timeout_s, worker_registry, maybe_origin_event_tx.clone(), has_peers, @@ -1207,6 +1470,7 @@ impl SimpleScheduler { task_abandoned_sweep_spawn, task_fleet_exchange_spawn, worker_match_logging_interval, + typed_properties: Mutex::new(HashMap::new()), unsatisfiable_tracker: Mutex::new(UnsatisfiableTracker::new( unsatisfiable_action_timeout, )), @@ -1262,6 +1526,28 @@ impl WorkerScheduler for SimpleScheduler { self.worker_scheduler.add_worker(worker).await } + async fn worker_dispatch_accepted( + &self, + worker_id: &WorkerId, + operation_id: &OperationId, + ) -> Result<(), Error> { + self.worker_scheduler + .worker_dispatch_accepted(worker_id, operation_id) + .await + } + + async fn worker_dispatch_declined( + &self, + worker_id: &WorkerId, + operation_id: &OperationId, + reason: String, + needs_kb: Option, + ) -> Result<(), Error> { + self.worker_scheduler + .worker_dispatch_declined(worker_id, operation_id, reason, needs_kb) + .await + } + async fn update_action( &self, worker_id: &WorkerId, @@ -1284,6 +1570,16 @@ impl WorkerScheduler for SimpleScheduler { .await } + async fn worker_liveness_refreshed( + &self, + worker_id: &WorkerId, + timestamp: WorkerTimestamp, + ) -> Result<(), Error> { + self.worker_scheduler + .worker_liveness_refreshed(worker_id, timestamp) + .await + } + async fn remove_worker(&self, worker_id: &WorkerId) -> Result<(), Error> { self.worker_scheduler.remove_worker(worker_id).await } @@ -1308,6 +1604,10 @@ impl WorkerScheduler for SimpleScheduler { .await } + async fn worker_snapshot(&self) -> Vec { + self.worker_scheduler.worker_snapshot().await + } + async fn kill_revoked_operations(&self) -> Result<(), Error> { self.worker_scheduler.kill_revoked_operations().await } diff --git a/nativelink-scheduler/src/simple_scheduler_state_manager.rs b/nativelink-scheduler/src/simple_scheduler_state_manager.rs index f75e9d889f1..664aa1eba4c 100644 --- a/nativelink-scheduler/src/simple_scheduler_state_manager.rs +++ b/nativelink-scheduler/src/simple_scheduler_state_manager.rs @@ -16,6 +16,7 @@ use core::ops::Bound; use core::time::Duration; use std::string::ToString; use std::sync::{Arc, Weak}; +use std::time::SystemTime; use async_lock::Mutex; use async_trait::async_trait; @@ -59,7 +60,7 @@ const MAX_RETRY_JITTER_MS: u64 = 20; /// lost a version conflict: exponential in the attempt, plus jitter. fn version_conflict_backoff(retry_count: usize) -> Duration { let base_delay = BASE_RETRY_DELAY_MS * (1 << retry_count.saturating_sub(2).min(4)); - let jitter = std::time::SystemTime::now() + let jitter = SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .map_or(0, |d| { u64::try_from(d.as_nanos()).expect("u64 error") % MAX_RETRY_JITTER_MS @@ -210,6 +211,10 @@ where async fn as_action_info(&self) -> Result<(Arc, Option), Error> { self.inner.as_action_info().await } + + async fn client_last_seen(&self) -> Result, Error> { + self.inner.client_last_seen().await + } } struct MatchingEngineActionStateResult @@ -266,6 +271,15 @@ where )) } + async fn client_last_seen(&self) -> Result, Error> { + let awaited_action = self + .awaited_action_sub + .borrow() + .await + .err_tip(|| "In MatchingEngineActionStateResult::client_last_seen")?; + Ok(Some(awaited_action.last_client_keepalive_timestamp())) + } + async fn changed(&mut self) -> Result<(Arc, Option), Error> { let mut timeout_attempts = 0; loop { @@ -350,10 +364,42 @@ where } } +/// How often an action is queued again, by cause, before the scheduler +/// completes it with `FailedPrecondition`. Each cause has its own budget: +/// an action's own failures, the workers it lost through no fault of its +/// own, and the memory escalations it was granted. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RetryLimits { + /// Internal errors and timeouts the action produced. + pub max_job_retries: usize, + /// Workers lost while the action ran. + pub max_worker_loss_retries: usize, + /// Memory escalations; 0 leaves only the fleet's ceiling to bound them. + pub max_escalations: usize, +} + +impl RetryLimits { + /// Lost workers allowed by default; the action did nothing wrong. + pub const DEFAULT_WORKER_LOSS_RETRIES: usize = 10; +} + +/// Only the job budget stated: the other two take their defaults, which is +/// what older callers meant. +impl From for RetryLimits { + fn from(max_job_retries: usize) -> Self { + Self { + max_job_retries, + max_worker_loss_retries: Self::DEFAULT_WORKER_LOSS_RETRIES, + max_escalations: 0, + } + } +} + /// `SimpleSchedulerStateManager` is responsible for maintaining the state of the scheduler. /// Scheduler state includes the actions that are queued, active, and recently completed. /// It also includes the workers that are available to execute actions based on allocation /// strategy. + #[derive(MetricsComponent, Debug)] pub struct SimpleSchedulerStateManager where @@ -371,6 +417,14 @@ where #[metric(help = "Maximum number of times a job can be retried")] max_job_retries: usize, + #[metric(help = "Maximum number of times a job is queued again after losing its worker")] + max_worker_loss_retries: usize, + + #[metric( + help = "Maximum number of memory escalations for a job; 0 leaves the ceiling to bound them" + )] + max_escalations: usize, + /// Duration after which an action is considered to be timed out if /// no event is received. #[metric( @@ -411,7 +465,7 @@ where NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static, { pub fn new( - max_job_retries: usize, + retry_limits: impl Into, no_event_action_timeout: Duration, client_action_timeout: Duration, max_executing_timeout: Duration, @@ -419,9 +473,16 @@ where now_fn: NowFn, worker_registry: Option, ) -> Arc { + let RetryLimits { + max_job_retries, + max_worker_loss_retries, + max_escalations, + } = retry_limits.into(); Arc::new_cyclic(|weak_self| Self { action_db, max_job_retries, + max_worker_loss_retries, + max_escalations, no_event_action_timeout, client_action_timeout, max_executing_timeout, @@ -913,6 +974,8 @@ where UpdateOperationType::UpdateWithError(_) => "Error", UpdateOperationType::UpdateWithDisconnect => "Disconnect", UpdateOperationType::ExecutionComplete => "ExecutionComplete", + UpdateOperationType::UpdateWithDecline(_) => "Decline", + UpdateOperationType::UpdateWithEscalation(_) => "Escalation", }; debug!( @@ -1065,15 +1128,15 @@ where } } UpdateOperationType::UpdateWithDisconnect => { - // A worker disconnect (e.g. OOMKill, pod eviction, network - // drop) used to requeue without counting as an attempt, - // which let an action that always crashes its worker loop - // forever until the Bazel client's --test_timeout fired. - // Count disconnects as attempts so max_job_retries caps the - // loop and the client sees a backend-attributable error. - awaited_action.attempts += 1; - - if awaited_action.attempts > self.max_job_retries { + // A lost worker (OOM-killed as a whole, pod evicted, + // network drop, timed out) is usually not this action's + // doing, so it has its own budget rather than spending + // the action's `max_job_retries`. The budget is what + // stops an action that takes a worker down every time + // it runs from doing so forever. + awaited_action.worker_losses += 1; + + if awaited_action.worker_losses > self.max_worker_loss_retries { ActionStage::Completed(ActionResult { execution_metadata: ExecutionMetadata { worker: maybe_worker_id @@ -1082,9 +1145,9 @@ where }, error: Some(make_err!( Code::FailedPrecondition, - "Worker disconnected repeatedly while executing this action ({} > {} attempts); the worker was likely OOM-killed or its pod evicted. Give the action a memory reservation or raise the pool's memory limit before retrying. {}", - awaited_action.attempts, - self.max_job_retries, + "The worker running this action was lost {} times, more than max_worker_loss_retries ({}); the action is likely what takes the worker down (OOM-killed pod, eviction). Give it a memory reservation or raise the pool's memory limit before retrying. {}", + awaited_action.worker_losses, + self.max_worker_loss_retries, format!( "for operation_id: {operation_id}, maybe_worker_id: {maybe_worker_id:?}" ), @@ -1096,7 +1159,62 @@ where ActionStage::Queued } } + UpdateOperationType::UpdateWithEscalation(escalation) => { + // The scheduler's own sizing decision, not a failure of + // the action: it has its own budget and does not spend + // `max_job_retries`. The fleet's ceiling bounds it + // anyway; the budget catches a fleet whose ceiling is + // far above where the action started. + awaited_action.escalations += 1; + if self.max_escalations > 0 && awaited_action.escalations > self.max_escalations + { + ActionStage::Completed(ActionResult { + execution_metadata: ExecutionMetadata { + worker: maybe_worker_id + .map_or_else(String::default, ToString::to_string), + ..ExecutionMetadata::default() + }, + error: Some( + make_err!( + Code::FailedPrecondition, + "Job cancelled after being killed for memory at every reservation tried ({} escalations, more than max_steps {}); the last reservation was {} {} {}", + awaited_action.escalations, + self.max_escalations, + escalation.value, + escalation.property, + format!("for operation_id: {operation_id}, maybe_worker_id: {maybe_worker_id:?}"), + ) + .merge(escalation.reason.clone()), + ), + ..ActionResult::default() + }) + } else { + // The same action, asking for more: the next + // dispatch carries the raised reservation and the + // worker holds it to that. The last step also asks + // for the worker's whole CPU, so it runs alone. + let mut action_info = (**awaited_action.action_info()).clone(); + action_info + .platform_properties + .insert(escalation.property.clone(), escalation.value.to_string()); + if let Some((cpu_property, cpu_value)) = &escalation.cpu { + action_info + .platform_properties + .insert(cpu_property.clone(), cpu_value.to_string()); + } + awaited_action.set_action_info(Arc::new(action_info)); + is_retry = true; + ActionStage::Queued + } + } // We shouldn't get here, but we just ignore it if we do. + UpdateOperationType::UpdateWithDecline(decline) => { + // Nothing ran, so nothing is charged against the retry + // cap and the retry metric does not move; the action + // goes back where it was. + debug!(reason = %decline.reason, "Dispatch declined, requeuing untried"); + ActionStage::Queued + } UpdateOperationType::ExecutionComplete => { warn!("inner_update_operation got an ExecutionComplete, that's unexpected."); return Ok(()); @@ -1267,12 +1385,9 @@ where .get_all_awaited_actions() .await .err_tip(|| "In SimpleSchedulerStateManager::filter_operations")? - .and_then(|awaited_action_subscriber| async move { - let awaited_action = awaited_action_subscriber - .borrow() - .await - .err_tip(|| "In SimpleSchedulerStateManager::filter_operations")?; - Ok((awaited_action_subscriber, awaited_action)) + .filter_map(|listed| async move { listed_or_lost(listed) }) + .try_filter_map(|awaited_action_subscriber| async move { + borrow_listed(awaited_action_subscriber).await }) .try_filter_map(|(subscriber, awaited_action)| { let filter = filter.clone(); @@ -1314,12 +1429,9 @@ where ) .await .err_tip(|| "In SimpleSchedulerStateManager::filter_operations")? - .and_then(|awaited_action_subscriber| async move { - let awaited_action = awaited_action_subscriber - .borrow() - .await - .err_tip(|| "In SimpleSchedulerStateManager::filter_operations")?; - Ok((awaited_action_subscriber, awaited_action)) + .filter_map(|listed| async move { listed_or_lost(listed) }) + .try_filter_map(|awaited_action_subscriber| async move { + borrow_listed(awaited_action_subscriber).await }) .try_filter_map(move |(subscriber, awaited_action)| { let filter = filter.clone(); @@ -1340,6 +1452,49 @@ where } } +/// A listed entry the store could not even hand over: the row was gone +/// (`NotFound`) or would not decode (`InvalidArgument`), which the Redis +/// listing reports as an error item before any subscriber exists. Skipped +/// and counted like a record that fails to read below. Anything else is the +/// store failing and stays an error for the reader. +fn listed_or_lost(listed: Result) -> Option> { + match listed { + Err(err) if is_lost_record(&err) => { + warn!( + ?err, + "Operation listed but its row cannot be read; skipping it in the listing" + ); + record_awaited_action_orphan("listing"); + None + } + listed => Some(listed), + } +} + +/// Reads the record behind a listed subscriber. A record the index listed +/// but the store no longer has (it expired between the search and the +/// read, or eviction took it) is skipped and counted, not returned as an +/// error: one such row must not end a listing that the provisioner and the +/// admin API read every few seconds. Any other error is the store failing +/// and is passed on. The sweep and the matcher skip the same way. +async fn borrow_listed( + subscriber: S, +) -> Result, Error> { + match subscriber.borrow().await { + Ok(awaited_action) => Ok(Some((subscriber, awaited_action))), + Err(err) if is_lost_record(&err) => { + warn!( + ?err, + "Operation listed but its record cannot be read; skipping it in the listing" + ); + record_awaited_action_orphan("listing"); + Ok(None) + } + // The store itself failed; that is the listing's error to report. + Err(err) => Err(err).err_tip(|| "In SimpleSchedulerStateManager::filter_operations"), + } +} + #[async_trait] impl ClientStateManager for SimpleSchedulerStateManager where diff --git a/nativelink-scheduler/src/store_awaited_action_db.rs b/nativelink-scheduler/src/store_awaited_action_db.rs index 14ab52308d6..41897e0ad45 100644 --- a/nativelink-scheduler/src/store_awaited_action_db.rs +++ b/nativelink-scheduler/src/store_awaited_action_db.rs @@ -39,7 +39,7 @@ use nativelink_util::store_trait::{ use nativelink_util::task::JoinHandleDropGuard; use opentelemetry::KeyValue; use serde::{Deserialize, Serialize}; -use tokio::sync::Notify; +use tokio::sync::{Notify, OnceCell}; use tracing::{error, warn}; use crate::awaited_action_db::{ @@ -91,6 +91,15 @@ pub struct OperationSubscriber, last_known_keepalive_ts: AtomicU64, now_fn: NowFn, + /// The record as a listing read it, when this subscriber came from one. + /// A listing already loads every record it returns, so `borrow` serves + /// this instead of reading the record again; only the client keepalive, + /// kept under its own key, is read, once. The matching pass borrows each + /// queued action two or three times per pass, so without this a pass + /// cost several store round trips per queued action. + listed: Option, + /// `listed` with the client keepalive merged in, filled on first borrow. + listed_with_keepalive: OnceCell, // If the SchedulerSubscriptionManager is not reliable, then this is populated // when the state is set to subscribed. When set it causes the state to be polled // as well as listening for the publishing. @@ -141,12 +150,20 @@ where last_known_keepalive_ts: AtomicU64::new(0), state: OperationSubscriberState::Unsubscribed, now_fn, + listed: None, + listed_with_keepalive: OnceCell::const_new(), maybe_last_stage: None, retain_completed_for, client_keepalive_ttl, } } + /// Serves `borrow` from the record a listing already read. + fn listed(mut self, awaited_action: AwaitedAction) -> Self { + self.listed = Some(awaited_action); + self + } + async fn inner_get_awaited_action( store: &S, key: OperationIdToAwaitedAction<'_>, @@ -166,14 +183,30 @@ where if let Some(client_operation_id) = maybe_client_operation_id { awaited_action.set_client_operation_id(client_operation_id); } + Self::merge_client_keepalive( + store, + key.0.as_ref(), + &mut awaited_action, + last_known_keepalive_ts, + ) + .await; + Ok(awaited_action) + } + /// Folds the client keepalive kept under its own key into the action's + /// timestamp and remembers it. + async fn merge_client_keepalive( + store: &S, + operation_id: &OperationId, + awaited_action: &mut AwaitedAction, + last_known_keepalive_ts: &AtomicU64, + ) { // Helper to convert SystemTime to unix timestamp let to_unix_ts = |t: SystemTime| -> u64 { t.duration_since(UNIX_EPOCH).map_or(0, |d| d.as_secs()) }; // Check the separate keepalive key for the most recent timestamp. let keepalive_ts = if USE_SEPARATE_CLIENT_KEEPALIVE_KEY { - let operation_id = key.0.as_ref(); match store.get_and_decode(ClientKeepaliveKey(operation_id)).await { Ok(Some(ts)) => { let awaited_ts = to_unix_ts(awaited_action.last_client_keepalive_timestamp()); @@ -192,7 +225,6 @@ where }; last_known_keepalive_ts.store(keepalive_ts, Ordering::Release); - Ok(awaited_action) } #[expect(clippy::future_not_send)] // TODO(jhpratt) remove this @@ -201,6 +233,23 @@ where .weak_store .upgrade() .err_tip(|| "Store gone in OperationSubscriber::get_awaited_action")?; + if let Some(listed) = &self.listed { + let merged = self + .listed_with_keepalive + .get_or_init(|| async { + let mut awaited_action = listed.clone(); + Self::merge_client_keepalive( + store.as_ref(), + self.subscription_key.0.as_ref(), + &mut awaited_action, + &self.last_known_keepalive_ts, + ) + .await; + awaited_action + }) + .await; + return Ok(merged.clone()); + } Self::inner_get_awaited_action( store.as_ref(), self.subscription_key.borrow(), @@ -1281,6 +1330,7 @@ where self.retain_completed_for, self.client_keepalive_ttl, ) + .listed(awaited_action) })) } @@ -1301,6 +1351,7 @@ where self.retain_completed_for, self.client_keepalive_ttl, ) + .listed(awaited_action) })) } } diff --git a/nativelink-scheduler/src/worker.rs b/nativelink-scheduler/src/worker.rs index 20da83cdc93..6770cdbaf12 100644 --- a/nativelink-scheduler/src/worker.rs +++ b/nativelink-scheduler/src/worker.rs @@ -20,14 +20,15 @@ use std::time::{SystemTime, UNIX_EPOCH}; use nativelink_error::{Code, Error, ResultExt}; use nativelink_metric::MetricsComponent; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{ - ConnectionResult, KillOperationRequest, StartExecute, UpdateForWorker, WorkerLoad, - update_for_worker, + ActionResourceUsage, ConnectionResult, KillOperationRequest, StartExecute, UpdateForWorker, + WorkerLoad, update_for_worker, }; use nativelink_util::action_messages::{ActionInfo, OperationId, WorkerId}; use nativelink_util::metrics_utils::{AsyncCounterWrapper, CounterWithTime, FuncCounterWrapper}; use nativelink_util::origin_event::OriginMetadata; use nativelink_util::platform_properties::{PlatformProperties, PlatformPropertyValue}; -use tokio::sync::mpsc::UnboundedSender; +use tokio::sync::mpsc::Sender; +use tokio::sync::mpsc::error::TrySendError; pub type WorkerTimestamp = u64; @@ -52,8 +53,9 @@ pub struct ActionInfoWithProps { /// Notifications to send worker about a requested state change. #[derive(Debug)] pub enum WorkerUpdate { - /// Requests that the worker begin executing this action. - RunAction(Box<(OperationId, ActionInfoWithProps)>), + /// Requests that the worker begin executing this action, dispatched at + /// this scheduler time. + RunAction(Box<(OperationId, ActionInfoWithProps, WorkerTimestamp)>), /// Request that the worker is no longer in the pool and may discard any jobs. Disconnect, @@ -74,6 +76,17 @@ pub struct PendingActionInfoData { /// worker's own bookkeeping. #[metric(help = "When the worker was asked to kill this operation.")] pub kill_requested_at: Option, + /// The worker's last-seen timestamp when the dispatch was sent. + #[metric(help = "When this operation was dispatched to the worker.")] + pub dispatched_at: WorkerTimestamp, + /// The worker said it took the action. A worker that does not speak + /// the acknowledgement never sets this, which is why the + /// unacknowledged sweep is opt-in. + #[metric(help = "Whether the worker acknowledged the dispatch.")] + pub accepted: bool, + /// The last resource usage the worker reported for this operation; it + /// arrives just before the result and says whether the result is a kill. + pub last_usage: Option, } /// Represents a connection to a worker and used as the medium to @@ -93,7 +106,7 @@ pub struct Worker { pub total_platform_properties: PlatformProperties, /// Channel to send commands from scheduler to worker. - pub tx: UnboundedSender, + pub tx: Sender, /// The action info of the running actions on the worker. #[metric(group = "running_action_infos")] @@ -109,6 +122,11 @@ pub struct Worker { #[metric(help = "If the worker is paused.")] pub is_paused: bool, + /// Set when the pause came from a decline for load: it lifts only on a + /// keepalive whose free memory covers this much, not on any keepalive. + #[metric(help = "Free memory in KiB the worker must report before it is unpaused.")] + pub pause_needs_kb: Option, + /// Whether the worker is draining. #[metric(help = "If the worker is draining.")] pub is_draining: bool, @@ -126,12 +144,40 @@ pub struct Worker { metrics: Arc, } +/// Messages the channel holds beyond the worker's own concurrency: kills, +/// keepalives and the disconnect, which arrive alongside the dispatches. +const CHANNEL_HEADROOM: usize = 16; +/// A worker that set no `max_inflight_tasks` gets this much queue. +const DEFAULT_CHANNEL_DISPATCHES: usize = 256; + +/// How many messages the worker's channel holds: what its concurrency +/// allows in flight plus the headroom. Bounded so that a worker that has +/// stopped reading cannot be handed work without limit; a full channel is +/// `ResourceExhausted` on the send, and the caller requeues. +pub fn channel_capacity(max_inflight_tasks: u64) -> usize { + let ceiling = DEFAULT_CHANNEL_DISPATCHES * 4; + let dispatches = match usize::try_from(max_inflight_tasks) { + Ok(0) => DEFAULT_CHANNEL_DISPATCHES, + Ok(n) if n <= ceiling => n, + _ => ceiling, + }; + dispatches + CHANNEL_HEADROOM +} + fn send_msg_to_worker( - tx: &UnboundedSender, + tx: &Sender, msg: update_for_worker::Update, ) -> Result<(), Error> { - tx.send(UpdateForWorker { update: Some(msg) }) - .map_err(|err| Error::from_std_err(Code::Internal, &err).append("Worker disconnected")) + match tx.try_send(UpdateForWorker { update: Some(msg) }) { + Ok(()) => Ok(()), + Err(TrySendError::Full(_)) => Err(Error::new( + Code::ResourceExhausted, + "Worker channel full".to_string(), + )), + Err(err @ TrySendError::Closed(_)) => { + Err(Error::from_std_err(Code::Internal, &err).append("Worker disconnected")) + } + } } /// Reduces the platform properties available on the worker based on the platform properties provided. @@ -158,7 +204,7 @@ impl Worker { pub fn new( id: WorkerId, platform_properties: PlatformProperties, - tx: UnboundedSender, + tx: Sender, timestamp: WorkerTimestamp, max_inflight_tasks: u64, ) -> Self { @@ -170,6 +216,7 @@ impl Worker { running_action_infos: HashMap::new(), last_update_timestamp: timestamp, is_paused: false, + pause_needs_kb: None, is_draining: false, max_inflight_tasks, last_load: None, @@ -183,17 +230,27 @@ impl Worker { keep_alive: FuncCounterWrapper::default(), notify_disconnect: CounterWithTime::default(), kill_operation: CounterWithTime::default(), + keep_alives_dropped: CounterWithTime::default(), }), } } /// Sends the initial connection information to the worker. This generally is just meta info. /// This should only be sent once and should always be the first item in the stream. - pub fn send_initial_connection_result(&mut self) -> Result<(), Error> { + /// `memory_property` is the platform property the scheduler vetoes + /// placement on; the worker declines for load against the same one, so + /// the two sides cannot disagree. None means the worker never declines + /// for load. + pub fn send_initial_connection_result( + &mut self, + memory_property: Option<&str>, + ) -> Result<(), Error> { send_msg_to_worker( &self.tx, update_for_worker::Update::ConnectionResult(ConnectionResult { worker_id: self.id.clone().into(), + dispatch_ack: true, + memory_property: memory_property.unwrap_or_default().to_string(), }), ) .err_tip(|| format!("Failed to send ConnectionResult to worker : {}", self.id)) @@ -203,8 +260,9 @@ impl Worker { pub async fn notify_update(&mut self, worker_update: WorkerUpdate) -> Result<(), Error> { match worker_update { WorkerUpdate::RunAction(action) => { - let (operation_id, action_info) = *action; - self.run_action(operation_id, action_info).await + let (operation_id, action_info, dispatched_at) = *action; + self.run_action(operation_id, action_info, dispatched_at) + .await } WorkerUpdate::Disconnect => { self.metrics.notify_disconnect.inc(); @@ -235,6 +293,28 @@ impl Worker { } } + /// Tells the worker to stop an operation this scheduler no longer holds + /// for it: a late acknowledgement of a dispatch the sweep already took + /// back means the worker is about to run an action that now belongs + /// elsewhere. Nothing to book, since the operation is not on the ledger. + pub(crate) fn kill_unknown_operation(&self, operation_id: &OperationId) -> Result<(), Error> { + self.metrics.kill_operation.inc(); + send_msg_to_worker( + &self.tx, + update_for_worker::Update::KillOperationRequest(KillOperationRequest { + operation_id: operation_id.to_string(), + }), + ) + } + + /// Forgets a kill request whose message never reached the worker, so + /// the next revoked-operation sweep sends it again. + pub(crate) fn clear_kill_request(&mut self, operation_id: &OperationId) { + if let Some(pending_action_info) = self.running_action_infos.get_mut(operation_id) { + pending_action_info.kill_requested_at = None; + } + } + /// Whether the worker has been told to kill this operation. pub(crate) fn is_kill_requested(&self, operation_id: &OperationId) -> bool { self.running_action_infos @@ -245,16 +325,47 @@ impl Worker { pub fn keep_alive(&mut self) -> Result<(), Error> { let tx = &mut self.tx; let id = &self.id; + let dropped = &self.metrics.keep_alives_dropped; self.metrics.keep_alive.wrap(move || { - send_msg_to_worker(tx, update_for_worker::Update::KeepAlive(())) - .err_tip(|| format!("Failed to send KeepAlive to worker : {id}")) + match send_msg_to_worker(tx, update_for_worker::Update::KeepAlive(())) { + // A worker that is not reading has a full queue of work + // ahead of the keepalive; skipping one costs nothing. + Err(err) if err.code == Code::ResourceExhausted => { + dropped.inc(); + Ok(()) + } + other => other.err_tip(|| format!("Failed to send KeepAlive to worker : {id}")), + } }) } + /// The worker said it took this operation. False when the operation is + /// not on this worker, which a late acknowledgement can cause. + pub fn mark_accepted(&mut self, operation_id: &OperationId) -> bool { + match self.running_action_infos.get_mut(operation_id) { + Some(pending) => { + pending.accepted = true; + true + } + None => false, + } + } + + /// Operations dispatched at or before `cutoff` that the worker never + /// acknowledged. + pub fn unacknowledged(&self, cutoff: WorkerTimestamp) -> Vec { + self.running_action_infos + .iter() + .filter(|(_, pending)| !pending.accepted && pending.dispatched_at <= cutoff) + .map(|(operation_id, _)| operation_id.clone()) + .collect() + } + async fn run_action( &mut self, operation_id: OperationId, action_info: ActionInfoWithProps, + dispatched_at: WorkerTimestamp, ) -> Result<(), Error> { let tx = &mut self.tx; let worker_platform_properties = &mut self.platform_properties; @@ -273,6 +384,9 @@ impl Worker { worker_id, request_metadata: action_info.origin_metadata.bazel_metadata.clone(), }; + // Send first: a channel that will not take the message + // leaves the ledger untouched, and the caller requeues. + send_msg_to_worker(tx, update_for_worker::Update::StartAction(start_execute))?; reduce_platform_properties( worker_platform_properties, &action_info.platform_properties, @@ -282,10 +396,12 @@ impl Worker { PendingActionInfoData { action_info, kill_requested_at: None, + dispatched_at, + accepted: false, + last_usage: None, }, ); - - send_msg_to_worker(tx, update_for_worker::Update::StartAction(start_execute)) + Ok(()) }) .await } @@ -305,6 +421,7 @@ impl Worker { })?; self.restore_platform_properties(&pending_action_info.action_info.platform_properties); self.is_paused = false; + self.pause_needs_kb = None; self.metrics.actions_completed.inc(); Ok(()) } @@ -363,4 +480,6 @@ struct Metrics { notify_disconnect: CounterWithTime, #[metric(help = "The number of kill_operation sent to this worker.")] kill_operation: CounterWithTime, + #[metric(help = "Keepalives not sent because the worker's channel was full.")] + keep_alives_dropped: CounterWithTime, } diff --git a/nativelink-scheduler/src/worker_scheduler.rs b/nativelink-scheduler/src/worker_scheduler.rs index 7305488807c..8d0208b24a4 100644 --- a/nativelink-scheduler/src/worker_scheduler.rs +++ b/nativelink-scheduler/src/worker_scheduler.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; + use async_trait::async_trait; use nativelink_error::Error; use nativelink_metric::RootMetricsComponent; @@ -27,6 +29,25 @@ use crate::worker::{Worker, WorkerTimestamp}; /// `WorkerScheduler` interface is responsible for interactions between the scheduler /// and worker related operations. +/// One connected worker, as the admin API reports it. Property values are +/// strings on both maps, the way they were registered. +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub struct WorkerSummary { + pub id: String, + pub running_actions: u32, + pub max_inflight_tasks: u64, + pub is_paused: bool, + pub is_draining: bool, + /// Seconds since the epoch of the worker's last message. + pub last_update_timestamp: u64, + /// What the worker registered with. + pub platform_properties: HashMap, + /// What is left after the running actions' reservations. + pub available_platform_properties: HashMap, + /// What the worker last reported having to spare, if it reports. + pub free_memory_kb: Option, +} + #[async_trait] pub trait WorkerScheduler: Sync + Send + Unpin + RootMetricsComponent + 'static { /// Returns the platform property manager. @@ -35,6 +56,24 @@ pub trait WorkerScheduler: Sync + Send + Unpin + RootMetricsComponent + 'static /// Adds a worker to the scheduler and begin using it to execute actions (when able). async fn add_worker(&self, worker: Worker) -> Result<(), Error>; + /// The worker took the dispatched operation. + async fn worker_dispatch_accepted( + &self, + worker_id: &WorkerId, + operation_id: &OperationId, + ) -> Result<(), Error>; + + /// The worker will not run the dispatched operation; `reason` is the + /// worker's word for why, `needs_kb` what it said the action wanted + /// when the reason was load. + async fn worker_dispatch_declined( + &self, + worker_id: &WorkerId, + operation_id: &OperationId, + reason: String, + needs_kb: Option, + ) -> Result<(), Error>; + /// Updates the status of an action to the scheduler from the worker. async fn update_action( &self, @@ -62,6 +101,16 @@ pub trait WorkerScheduler: Sync + Send + Unpin + RootMetricsComponent + 'static load: Option, ) -> Result<(), Error>; + /// Any other message from the worker (an acknowledgement, a decline, an + /// execute result) proves it is alive. Unlike a keepalive it lifts no + /// pause: a decline is not the worker saying it is ready to be asked + /// again. + async fn worker_liveness_refreshed( + &self, + worker_id: &WorkerId, + timestamp: WorkerTimestamp, + ) -> Result<(), Error>; + /// Removes worker from pool and reschedule any tasks that might be running on it. async fn remove_worker(&self, worker_id: &WorkerId) -> Result<(), Error>; @@ -80,6 +129,10 @@ pub trait WorkerScheduler: Sync + Send + Unpin + RootMetricsComponent + 'static /// Sets if the worker is draining or not. async fn set_drain_worker(&self, worker_id: &WorkerId, is_draining: bool) -> Result<(), Error>; + /// Every connected worker as the scheduler sees it right now, for the + /// admin API: what it advertised, what it has left, what it runs. + async fn worker_snapshot(&self) -> Vec; + /// Tells workers to kill operations they are still running but the /// scheduler has finished, requeued or dropped without them (client /// timeouts and cancellations, execution deadlines, retries elsewhere). diff --git a/nativelink-scheduler/tests/admin_test.rs b/nativelink-scheduler/tests/admin_test.rs new file mode 100644 index 00000000000..9e621fa848c --- /dev/null +++ b/nativelink-scheduler/tests/admin_test.rs @@ -0,0 +1,307 @@ +use core::time::Duration; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::UNIX_EPOCH; + +use async_trait::async_trait; +use futures::stream; +use mock_instant::thread_local::MockClock; +use nativelink_config::schedulers::{PropertyType, SimpleSpec}; +use nativelink_error::{Code, Error, ResultExt, make_err}; +use nativelink_macro::nativelink_test; +use nativelink_metric::{ + MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent, +}; +use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::UpdateForWorker; +use nativelink_scheduler::admin::{QueuedDemand, queued_demand, queued_demand_json, workers_json}; +use nativelink_scheduler::default_scheduler_factory::memory_awaited_action_db_factory; +use nativelink_scheduler::simple_scheduler::SimpleScheduler; +use nativelink_scheduler::worker::Worker; +use nativelink_scheduler::worker_scheduler::{WorkerScheduler, WorkerSummary}; +use nativelink_util::action_messages::{ + ActionInfo, ActionStage, ActionState, OperationId, WorkerId, +}; +use nativelink_util::common::DigestInfo; +use nativelink_util::instant_wrapper::MockInstantWrapped; +use nativelink_util::operation_state_manager::{ + ActionStateResult, ActionStateResultStream, ClientStateManager, OperationFilter, +}; +use nativelink_util::origin_event::OriginMetadata; +use nativelink_util::platform_properties::{PlatformProperties, PlatformPropertyValue}; +use pretty_assertions::assert_eq; +use tokio::sync::{Notify, mpsc}; +use utils::scheduler_utils::make_base_action_info; + +mod utils { + pub(crate) mod scheduler_utils; +} + +const NOW_TIME: u64 = 10000; + +fn make_scheduler() -> Arc { + MockClock::set_time(Duration::from_secs(NOW_TIME)); + let task_change_notify = Arc::new(Notify::new()); + let (scheduler, _worker_scheduler) = SimpleScheduler::new_with_callback( + &SimpleSpec { + supported_platform_properties: Some(HashMap::from([( + "cpu_count".to_string(), + PropertyType::Minimum, + )])), + client_action_timeout_s: 1_000_000, + ..SimpleSpec::default() + }, + memory_awaited_action_db_factory(0, &task_change_notify, MockInstantWrapped::default), + || async move {}, + task_change_notify, + MockInstantWrapped::default, + None, + ); + scheduler +} + +async fn add_worker( + scheduler: &SimpleScheduler, + name: &str, + cpu_count: u64, +) -> Result, Error> { + let (tx, mut rx) = mpsc::channel(64); + let worker = Worker::new( + WorkerId(name.to_string()), + PlatformProperties::new(HashMap::from([( + "cpu_count".to_string(), + PlatformPropertyValue::Minimum(cpu_count), + )])), + tx, + NOW_TIME, + 4, + ); + scheduler + .add_worker(worker) + .await + .err_tip(|| "Failed to add worker")?; + tokio::task::yield_now().await; + rx.recv().await.expect("the connection result"); + Ok(rx) +} + +async fn add_action( + scheduler: &SimpleScheduler, + digest_byte: u8, + priority: i32, + cpu_count: &str, +) -> Result, Error> { + let mut action_info = make_base_action_info( + UNIX_EPOCH + MockClock::time(), + DigestInfo::new([digest_byte; 32], 512), + ); + let inner = Arc::make_mut(&mut action_info); + inner.priority = priority; + inner.platform_properties = HashMap::from([("cpu_count".to_string(), cpu_count.to_string())]); + let result = scheduler + .add_action(OperationId::default(), action_info) + .await?; + tokio::task::yield_now().await; + Ok(result) +} + +/// The demand listing is what a provisioner scales on: every queued action, +/// highest priority first, with what it asked for and when its client was +/// last heard from. +#[nativelink_test] +async fn queued_demand_lists_waiting_actions_highest_priority_first() -> Result<(), Error> { + let scheduler = make_scheduler(); + let _low = add_action(&scheduler, 1, 0, "2").await?; + let _high = add_action(&scheduler, 2, 5, "1").await?; + + let demand = queued_demand(scheduler.as_ref()).await?; + assert_eq!(demand.len(), 2); + assert_eq!(demand[0].priority, 5); + assert_eq!( + demand[0].platform_properties, + HashMap::from([("cpu_count".to_string(), "1".to_string())]) + ); + assert_eq!(demand[1].priority, 0); + for entry in &demand { + assert_eq!(entry.queued_since_ms, NOW_TIME * 1000); + let last_seen = entry + .client_last_seen_ms + .expect("a client is still waiting on the action"); + assert!(last_seen >= entry.queued_since_ms); + assert_ne!(entry.operation_id, ""); + } + Ok(()) +} + +/// An action a worker is running is no longer demand, and the worker +/// listing shows it running with the room it took. +#[nativelink_test] +async fn running_actions_leave_the_demand_and_show_on_the_worker() -> Result<(), Error> { + let scheduler = make_scheduler(); + let _rx = add_worker(&scheduler, "w1", 4).await?; + let action = add_action(&scheduler, 1, 0, "3").await?; + scheduler.do_try_match_for_test().await?; + assert_eq!(action.as_state().await?.0.stage, ActionStage::Executing); + + assert_eq!(queued_demand(scheduler.as_ref()).await?.len(), 0); + + let workers = scheduler.worker_snapshot().await; + assert_eq!(workers.len(), 1); + let worker = &workers[0]; + assert_eq!(worker.id, "w1"); + assert_eq!(worker.running_actions, 1); + assert_eq!(worker.max_inflight_tasks, 4); + assert!(!worker.is_paused && !worker.is_draining); + assert_eq!(worker.last_update_timestamp, NOW_TIME); + assert_eq!(worker.platform_properties["cpu_count"], "4"); + assert_eq!(worker.available_platform_properties["cpu_count"], "1"); + assert_eq!(worker.free_memory_kb, None); + Ok(()) +} + +/// Both listings serialize to JSON the way the admin endpoint serves them, +/// and read back to the same values. +#[nativelink_test] +async fn listings_round_trip_through_json() -> Result<(), Error> { + let scheduler = make_scheduler(); + let _rx = add_worker(&scheduler, "w1", 4).await?; + let _waiting = add_action(&scheduler, 1, 2, "8").await?; + scheduler.do_try_match_for_test().await?; + + let demand = queued_demand(scheduler.as_ref()).await?; + let parsed: Vec = + serde_json::from_str(&queued_demand_json(scheduler.as_ref()).await?) + .expect("demand JSON parses"); + assert_eq!(parsed, demand); + assert_eq!( + parsed.len(), + 1, + "an 8-cpu action does not fit a 4-cpu worker" + ); + + let workers = scheduler.worker_snapshot().await; + let parsed: Vec = + serde_json::from_str(&workers_json(&workers)?).expect("worker JSON parses"); + assert_eq!(parsed, workers); + Ok(()) +} + +/// One listed operation as the demand listing reads it back: the filter +/// saw it queued, but each later read is fresh and may find it moved on, +/// gone, or unreadable. +#[derive(Clone)] +struct Listed { + stage: Result, +} + +#[async_trait] +impl ActionStateResult for Listed { + async fn as_state(&self) -> Result<(Arc, Option), Error> { + let stage = self + .stage + .clone() + .map_err(|code| make_err!(code, "record read failed"))?; + Ok(( + Arc::new(ActionState { + stage, + client_operation_id: OperationId::default(), + action_digest: DigestInfo::zero_digest(), + last_transition_timestamp: UNIX_EPOCH, + }), + None, + )) + } + + async fn changed(&mut self) -> Result<(Arc, Option), Error> { + self.as_state().await + } + + async fn as_action_info(&self) -> Result<(Arc, Option), Error> { + Ok(( + make_base_action_info(UNIX_EPOCH, DigestInfo::zero_digest()), + None, + )) + } +} + +/// A scheduler whose queued listing is exactly these operations. +struct FixedListing(Vec); + +impl MetricsComponent for FixedListing { + fn publish( + &self, + _kind: MetricKind, + _field_metadata: MetricFieldData, + ) -> Result { + Ok(MetricPublishKnownKindData::Component) + } +} + +#[async_trait] +impl ClientStateManager for FixedListing { + async fn add_action( + &self, + _client_operation_id: OperationId, + _action_info: Arc, + ) -> Result, Error> { + Err(make_err!(Code::Unimplemented, "not part of this fixture")) + } + + async fn filter_operations( + &self, + _filter: OperationFilter, + ) -> Result { + let listed: Vec> = self + .0 + .iter() + .cloned() + .map(|listed| -> Box { Box::new(listed) }) + .collect(); + Ok(Box::pin(stream::iter(listed))) + } +} + +/// Between the filter and the read a record can be dispatched, expire or +/// turn out unreadable. Demand keeps only what is still queued, and one bad +/// record does not take the listing down with it. +#[nativelink_test] +async fn demand_leaves_out_what_moved_on_or_went_away() -> Result<(), Error> { + let scheduler = FixedListing(vec![ + Listed { + stage: Ok(ActionStage::Queued), + }, + Listed { + stage: Err(Code::NotFound), + }, + Listed { + stage: Ok(ActionStage::Executing), + }, + Listed { + stage: Err(Code::InvalidArgument), + }, + Listed { + stage: Ok(ActionStage::Queued), + }, + ]); + let demand = queued_demand(&scheduler).await?; + assert_eq!(demand.len(), 2, "only the two still queued are demand"); + Ok(()) +} + +/// A read that fails because the store is unreachable is the listing's +/// error, not a record to leave out. +#[nativelink_test] +async fn demand_fails_when_the_store_does() -> Result<(), Error> { + let scheduler = FixedListing(vec![ + Listed { + stage: Ok(ActionStage::Queued), + }, + Listed { + stage: Err(Code::Unavailable), + }, + ]); + let err = queued_demand(&scheduler) + .await + .expect_err("a store failure fails the listing"); + assert_eq!(err.code, Code::Unavailable); + Ok(()) +} diff --git a/nativelink-scheduler/tests/backpressure_test.rs b/nativelink-scheduler/tests/backpressure_test.rs index 63010c5ff01..032d09f2f4f 100644 --- a/nativelink-scheduler/tests/backpressure_test.rs +++ b/nativelink-scheduler/tests/backpressure_test.rs @@ -60,8 +60,8 @@ fn make_scheduler() -> Arc { async fn add_worker( scheduler: &SimpleScheduler, name: &str, -) -> Result, Error> { - let (tx, mut rx) = mpsc::unbounded_channel(); +) -> Result, Error> { + let (tx, mut rx) = mpsc::channel(64); let worker = Worker::new( WorkerId(name.to_string()), PlatformProperties::new(HashMap::new()), @@ -95,7 +95,7 @@ async fn add_action(scheduler: &SimpleScheduler, digest_byte: u8) -> Result<(), } /// The operation ids of every `StartAction` the worker has been sent. -fn started(rx: &mut mpsc::UnboundedReceiver) -> Vec { +fn started(rx: &mut mpsc::Receiver) -> Vec { let mut ids = Vec::new(); while let Ok(msg) = rx.try_recv() { if let Some(update_for_worker::Update::StartAction(start)) = msg.update { diff --git a/nativelink-scheduler/tests/dispatch_ack_test.rs b/nativelink-scheduler/tests/dispatch_ack_test.rs new file mode 100644 index 00000000000..2e3c2020bc4 --- /dev/null +++ b/nativelink-scheduler/tests/dispatch_ack_test.rs @@ -0,0 +1,415 @@ +// Copyright 2024 The NativeLink Authors. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Dispatch is a proposal until the worker answers it: a decline sends the +//! action back untried and pauses the worker, a full channel does the same +//! without an eviction, and an acknowledgement that never comes is a +//! decline after the timeout. + +use core::time::Duration; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::UNIX_EPOCH; + +use mock_instant::thread_local::MockClock; +use nativelink_config::schedulers::{PropertyType, SimpleSpec}; +use nativelink_error::{Error, ResultExt}; +use nativelink_macro::nativelink_test; +use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{ + StartExecute, UpdateForWorker, WorkerLoad, update_for_worker, +}; +use nativelink_scheduler::default_scheduler_factory::memory_awaited_action_db_factory; +use nativelink_scheduler::simple_scheduler::SimpleScheduler; +use nativelink_scheduler::worker::Worker; +use nativelink_scheduler::worker_scheduler::WorkerScheduler; +use nativelink_util::action_messages::{ + ActionInfo, ActionResult, ActionStage, OperationId, WorkerId, +}; +use nativelink_util::common::DigestInfo; +use nativelink_util::instant_wrapper::MockInstantWrapped; +use nativelink_util::operation_state_manager::{ + ActionStateResult, ClientStateManager, UpdateOperationType, +}; +use nativelink_util::platform_properties::{PlatformProperties, PlatformPropertyValue}; +use pretty_assertions::assert_eq; +use tokio::sync::{Notify, mpsc}; +use utils::scheduler_utils::make_base_action_info; + +mod utils { + pub(crate) mod scheduler_utils; +} + +const NOW_TIME: u64 = 10000; +const WORKER: &str = "the_worker"; + +fn make_scheduler( + max_job_retries: usize, + dispatch_ack_timeout_s: u64, +) -> (Arc, Arc) { + MockClock::set_time(Duration::from_secs(NOW_TIME)); + let task_change_notify = Arc::new(Notify::new()); + let spec = SimpleSpec { + supported_platform_properties: Some(HashMap::from([( + "memory_kb".to_string(), + PropertyType::Minimum, + )])), + max_job_retries, + dispatch_ack_timeout_s, + ..SimpleSpec::default() + }; + SimpleScheduler::new_with_callback( + &spec, + memory_awaited_action_db_factory(0, &task_change_notify, MockInstantWrapped::default), + || async move {}, + task_change_notify, + MockInstantWrapped::default, + None, + ) +} + +fn worker_id() -> WorkerId { + WorkerId(WORKER.to_string()) +} + +fn worker_with_channel(capacity: usize) -> (Worker, mpsc::Receiver) { + let (tx, rx) = mpsc::channel(capacity); + let worker = Worker::new( + worker_id(), + PlatformProperties::new(HashMap::from([( + "memory_kb".to_string(), + PlatformPropertyValue::Minimum(5_000), + )])), + tx, + NOW_TIME, + /* max_inflight_tasks */ 4, + ); + (worker, rx) +} + +async fn add_worker(scheduler: &SimpleScheduler) -> Result, Error> { + let (worker, mut rx) = worker_with_channel(64); + scheduler + .add_worker(worker) + .await + .err_tip(|| "Failed to add worker")?; + tokio::task::yield_now().await; + let connected = rx.recv().await.unwrap(); + let Some(update_for_worker::Update::ConnectionResult(result)) = connected.update else { + panic!("expected a ConnectionResult, got {connected:?}"); + }; + assert!( + result.dispatch_ack, + "the scheduler announces the acknowledgement" + ); + assert!( + result.memory_property.is_empty(), + "no live memory veto, so the worker is told nothing to decline for load on" + ); + Ok(rx) +} + +async fn add_action( + scheduler: &SimpleScheduler, + seed: u8, +) -> Result, Error> { + let base = make_base_action_info( + UNIX_EPOCH + MockClock::time(), + DigestInfo::new([seed; 32], 512), + ); + let action_info = Arc::new(ActionInfo { + platform_properties: HashMap::from([("memory_kb".to_string(), "2000".to_string())]), + ..(*base).clone() + }); + let listener = scheduler + .add_action(OperationId::default(), action_info) + .await?; + tokio::task::yield_now().await; + Ok(listener) +} + +async fn next_dispatch(rx: &mut mpsc::Receiver) -> StartExecute { + let msg = tokio::time::timeout(Duration::from_secs(2), rx.recv()) + .await + .expect("a dispatch within two seconds") + .expect("the channel is open"); + let Some(update_for_worker::Update::StartAction(start_execute)) = msg.update else { + panic!("expected a StartAction, got {msg:?}"); + }; + start_execute +} + +async fn no_dispatch(rx: &mut mpsc::Receiver) { + let quiet = tokio::time::timeout(Duration::from_millis(200), rx.recv()).await; + assert!( + quiet.is_err(), + "expected no dispatch, got {:?}", + quiet.expect("timeout branch handled above") + ); +} + +async fn keepalive( + worker_scheduler: &dyn WorkerScheduler, + at: u64, + free_memory_kb: Option, +) -> Result<(), Error> { + worker_scheduler + .worker_keep_alive_received( + &worker_id(), + at, + free_memory_kb.map(|free_memory_kb| WorkerLoad { free_memory_kb }), + ) + .await?; + tokio::task::yield_now().await; + Ok(()) +} + +async fn complete(worker_scheduler: &dyn WorkerScheduler, operation_id: &str) -> Result<(), Error> { + worker_scheduler + .update_action( + &worker_id(), + &OperationId::from(operation_id), + UpdateOperationType::UpdateWithActionStage(ActionStage::Completed( + ActionResult::default(), + )), + ) + .await?; + tokio::task::yield_now().await; + Ok(()) +} + +async fn wait_for_completion(listener: &mut Box) -> ActionResult { + loop { + let (state, _) = listener.changed().await.expect("state stream open"); + if let ActionStage::Completed(result) = &state.stage { + return result.clone(); + } + } +} + +/// A decline sends the action back untried: with a retry cap of one, two +/// declines and a completion still succeed. The worker is paused by the +/// decline; a decline for load lifts only on a keepalive whose free memory +/// covers the need, any other on the next keepalive. +#[nativelink_test] +async fn a_decline_requeues_untried_and_pauses_until_the_load_fits() -> Result<(), Error> { + let (scheduler, worker_scheduler) = make_scheduler(1, 0); + let mut rx = add_worker(&scheduler).await?; + let mut listener = add_action(&scheduler, 1).await?; + + let first = next_dispatch(&mut rx).await; + worker_scheduler + .worker_dispatch_declined( + &worker_id(), + &OperationId::from(first.operation_id.as_str()), + "load".to_string(), + Some(2_000), + ) + .await?; + tokio::task::yield_now().await; + no_dispatch(&mut rx).await; + + // Too little free memory: still paused. + keepalive(worker_scheduler.as_ref(), NOW_TIME + 1, Some(1_000)).await?; + no_dispatch(&mut rx).await; + + // Enough: the same operation comes back. + keepalive(worker_scheduler.as_ref(), NOW_TIME + 2, Some(3_000)).await?; + let second = next_dispatch(&mut rx).await; + assert_eq!(second.operation_id, first.operation_id); + + // A decline for capacity lifts on any keepalive. + worker_scheduler + .worker_dispatch_declined( + &worker_id(), + &OperationId::from(second.operation_id.as_str()), + "at_capacity".to_string(), + None, + ) + .await?; + tokio::task::yield_now().await; + no_dispatch(&mut rx).await; + keepalive(worker_scheduler.as_ref(), NOW_TIME + 3, None).await?; + let third = next_dispatch(&mut rx).await; + assert_eq!(third.operation_id, first.operation_id); + + complete(worker_scheduler.as_ref(), &third.operation_id).await?; + let result = wait_for_completion(&mut listener).await; + assert!( + result.error.is_none(), + "two declines did not count against a retry cap of one: {:?}", + result.error + ); + Ok(()) +} + +/// A worker whose channel will not take the dispatch keeps its place: the +/// action goes back untried, the worker is paused, and once it reads and +/// sends a keepalive the same action comes to it. +#[nativelink_test] +async fn a_full_channel_requeues_and_pauses_instead_of_evicting() -> Result<(), Error> { + let (scheduler, worker_scheduler) = make_scheduler(1, 0); + // Room for the ConnectionResult and nothing else. + let (worker, mut rx) = worker_with_channel(1); + scheduler.add_worker(worker).await?; + tokio::task::yield_now().await; + + let mut listener = add_action(&scheduler, 2).await?; + tokio::task::yield_now().await; + + // The worker was not evicted: draining it would fail otherwise. + worker_scheduler + .set_drain_worker(&worker_id(), false) + .await + .err_tip(|| "the worker is still in the pool")?; + + // The worker catches up: reads its connection result, sends a + // keepalive, and the action arrives. + let connected = rx.recv().await.unwrap(); + assert!(matches!( + connected.update, + Some(update_for_worker::Update::ConnectionResult(_)) + )); + keepalive(worker_scheduler.as_ref(), NOW_TIME + 1, None).await?; + let dispatched = next_dispatch(&mut rx).await; + complete(worker_scheduler.as_ref(), &dispatched.operation_id).await?; + let result = wait_for_completion(&mut listener).await; + assert!(result.error.is_none(), "{:?}", result.error); + Ok(()) +} + +/// With the sweep on, a dispatch the worker never acknowledged goes back +/// untried after the timeout; an acknowledged one stays. +#[nativelink_test] +async fn an_unacknowledged_dispatch_is_requeued_after_the_timeout() -> Result<(), Error> { + let (scheduler, worker_scheduler) = make_scheduler(1, 5); + let mut rx = add_worker(&scheduler).await?; + let mut listener = add_action(&scheduler, 3).await?; + + let first = next_dispatch(&mut rx).await; + // Within the timeout: nothing happens. + keepalive(worker_scheduler.as_ref(), NOW_TIME + 3, None).await?; + no_dispatch(&mut rx).await; + // Past it: requeued and the worker paused; the next keepalive brings + // the same operation back. + keepalive(worker_scheduler.as_ref(), NOW_TIME + 6, None).await?; + no_dispatch(&mut rx).await; + keepalive(worker_scheduler.as_ref(), NOW_TIME + 7, None).await?; + let second = next_dispatch(&mut rx).await; + assert_eq!(second.operation_id, first.operation_id); + + // Acknowledged this time: the sweep leaves it alone. + worker_scheduler + .worker_dispatch_accepted( + &worker_id(), + &OperationId::from(second.operation_id.as_str()), + ) + .await?; + keepalive(worker_scheduler.as_ref(), NOW_TIME + 20, None).await?; + no_dispatch(&mut rx).await; + + complete(worker_scheduler.as_ref(), &second.operation_id).await?; + let result = wait_for_completion(&mut listener).await; + assert!(result.error.is_none(), "{:?}", result.error); + Ok(()) +} + +/// A decline for an operation the worker no longer holds, because it +/// finished or the sweep already took it back, is nothing to act on: the +/// worker keeps taking work. +#[nativelink_test] +async fn a_late_decline_does_not_pause_the_worker() -> Result<(), Error> { + let (scheduler, worker_scheduler) = make_scheduler(1, 0); + let mut rx = add_worker(&scheduler).await?; + let mut listener = add_action(&scheduler, 1).await?; + let first = next_dispatch(&mut rx).await; + complete(worker_scheduler.as_ref(), &first.operation_id).await?; + let result = wait_for_completion(&mut listener).await; + assert!(result.error.is_none(), "{:?}", result.error); + + // The decline arrives after the completion. + worker_scheduler + .worker_dispatch_declined( + &worker_id(), + &OperationId::from(first.operation_id.as_str()), + "load".to_string(), + Some(5_000), + ) + .await?; + + let _listener = add_action(&scheduler, 2).await?; + next_dispatch(&mut rx).await; + Ok(()) +} + +async fn next_kill(rx: &mut mpsc::Receiver) -> String { + let msg = tokio::time::timeout(Duration::from_secs(2), rx.recv()) + .await + .expect("a kill within two seconds") + .expect("the channel is open"); + let Some(update_for_worker::Update::KillOperationRequest(kill)) = msg.update else { + panic!("expected a KillOperationRequest, got {msg:?}"); + }; + kill.operation_id +} + +/// The acknowledgement clock starts at the dispatch, not at the worker's +/// last message before it: a worker that was quiet for a while still gets +/// the whole timeout for a dispatch made later. +#[nativelink_test] +async fn the_acknowledgement_clock_starts_at_the_dispatch() -> Result<(), Error> { + let (scheduler, worker_scheduler) = make_scheduler(1, 5); + let mut rx = add_worker(&scheduler).await?; + // The worker's last message was at NOW_TIME; the dispatch is four + // seconds later. + MockClock::advance(Duration::from_secs(4)); + let _listener = add_action(&scheduler, 4).await?; + let first = next_dispatch(&mut rx).await; + + // Four seconds after the dispatch, eight after the worker's last + // message: still within the timeout, so the acknowledgement counts. + keepalive(worker_scheduler.as_ref(), NOW_TIME + 8, None).await?; + worker_scheduler + .worker_dispatch_accepted( + &worker_id(), + &OperationId::from(first.operation_id.as_str()), + ) + .await?; + keepalive(worker_scheduler.as_ref(), NOW_TIME + 20, None).await?; + no_dispatch(&mut rx).await; + Ok(()) +} + +/// An acknowledgement for a dispatch the sweep already took back means the +/// worker is about to run an action that now belongs elsewhere: it is told +/// to stop at once, not at the next revoked-operation sweep. +#[nativelink_test] +async fn a_late_acknowledgement_is_answered_with_a_kill() -> Result<(), Error> { + let (scheduler, worker_scheduler) = make_scheduler(1, 5); + let mut rx = add_worker(&scheduler).await?; + let _listener = add_action(&scheduler, 5).await?; + let first = next_dispatch(&mut rx).await; + + // Past the timeout with no acknowledgement: taken back, worker paused. + keepalive(worker_scheduler.as_ref(), NOW_TIME + 6, None).await?; + no_dispatch(&mut rx).await; + + worker_scheduler + .worker_dispatch_accepted( + &worker_id(), + &OperationId::from(first.operation_id.as_str()), + ) + .await?; + assert_eq!(next_kill(&mut rx).await, first.operation_id); + Ok(()) +} diff --git a/nativelink-scheduler/tests/escalation_test.rs b/nativelink-scheduler/tests/escalation_test.rs new file mode 100644 index 00000000000..1ebcb0a224b --- /dev/null +++ b/nativelink-scheduler/tests/escalation_test.rs @@ -0,0 +1,402 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use core::time::Duration; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::UNIX_EPOCH; + +use mock_instant::thread_local::MockClock; +use nativelink_config::schedulers::{MemoryEscalationSpec, PropertyType, SimpleSpec}; +use nativelink_error::{Code, Error, ResultExt, make_err}; +use nativelink_macro::nativelink_test; +use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{ + ActionResourceUsage, Reservation, ResourceOutcome, StartExecute, UpdateForWorker, + update_for_worker, +}; +use nativelink_scheduler::default_scheduler_factory::memory_awaited_action_db_factory; +use nativelink_scheduler::simple_scheduler::SimpleScheduler; +use nativelink_scheduler::worker::Worker; +use nativelink_scheduler::worker_scheduler::WorkerScheduler; +use nativelink_util::action_messages::{ + ActionInfo, ActionResult, ActionStage, OperationId, WorkerId, +}; +use nativelink_util::common::DigestInfo; +use nativelink_util::instant_wrapper::MockInstantWrapped; +use nativelink_util::operation_state_manager::{ + ActionStateResult, ClientStateManager, UpdateOperationType, +}; +use nativelink_util::platform_properties::{PlatformProperties, PlatformPropertyValue}; +use pretty_assertions::assert_eq; +use tokio::sync::{Notify, mpsc}; +use utils::scheduler_utils::make_base_action_info; + +mod utils { + pub(crate) mod scheduler_utils; +} + +const NOW_TIME: u64 = 10000; +const WORKER: &str = "the_worker"; + +fn make_scheduler(max_job_retries: usize) -> (Arc, Arc) { + make_scheduler_with_ladder(max_job_retries, vec![]) +} + +fn make_scheduler_with_ladder( + max_job_retries: usize, + ladder_kb: Vec, +) -> (Arc, Arc) { + make_scheduler_with_limits(max_job_retries, ladder_kb, 0) +} + +fn make_scheduler_with_limits( + max_job_retries: usize, + ladder_kb: Vec, + max_steps: u64, +) -> (Arc, Arc) { + MockClock::set_time(Duration::from_secs(NOW_TIME)); + let task_change_notify = Arc::new(Notify::new()); + let spec = SimpleSpec { + supported_platform_properties: Some(HashMap::from([( + "memory_kb".to_string(), + PropertyType::Minimum, + )])), + memory_escalation: Some(MemoryEscalationSpec { + property: "memory_kb".to_string(), + percent: 200, + max_kb: 0, + ladder_kb, + cpu_property: "cpu_count".to_string(), + max_steps, + }), + max_job_retries, + max_worker_loss_retries: 2, + ..SimpleSpec::default() + }; + SimpleScheduler::new_with_callback( + &spec, + memory_awaited_action_db_factory(0, &task_change_notify, MockInstantWrapped::default), + || async move {}, + task_change_notify, + MockInstantWrapped::default, + None, + ) +} + +async fn add_worker( + scheduler: &SimpleScheduler, + memory_kb: u64, +) -> Result, Error> { + let (tx, mut rx) = mpsc::channel(64); + let worker = Worker::new( + WorkerId(WORKER.to_string()), + PlatformProperties::new(HashMap::from([( + "memory_kb".to_string(), + PlatformPropertyValue::Minimum(memory_kb), + )])), + tx, + NOW_TIME, + /* max_inflight_tasks */ 4, + ); + scheduler + .add_worker(worker) + .await + .err_tip(|| "Failed to add worker")?; + tokio::task::yield_now().await; + let connected = rx.recv().await.unwrap(); + assert!(matches!( + connected.update, + Some(update_for_worker::Update::ConnectionResult(_)) + )); + Ok(rx) +} + +/// The next dispatch the worker sees, and the memory it was reserved with. +async fn next_dispatch(rx: &mut mpsc::Receiver) -> (StartExecute, u64) { + let msg = tokio::time::timeout(Duration::from_secs(2), rx.recv()) + .await + .expect("a dispatch within two seconds") + .expect("the channel is open"); + let Some(update_for_worker::Update::StartAction(start_execute)) = msg.update else { + panic!("expected a StartAction, got {msg:?}"); + }; + let memory_kb = start_execute + .platform + .as_ref() + .and_then(|platform| { + platform + .properties + .iter() + .find(|p| p.name == "memory_kb") + .and_then(|p| p.value.parse().ok()) + }) + .expect("the dispatch carries memory_kb"); + (start_execute, memory_kb) +} + +/// What a worker sends when its enforcement killed the action at its +/// reservation: the usage report, then the result carrying the error. +async fn report_memory_kill( + worker_scheduler: &dyn WorkerScheduler, + operation_id: &str, + reserved_kb: u64, +) -> Result<(), Error> { + let worker_id = WorkerId(WORKER.to_string()); + let operation_id = OperationId::from(operation_id); + worker_scheduler + .record_action_resource_usage( + &worker_id, + &operation_id, + ActionResourceUsage { + peak_memory_kb: reserved_kb + reserved_kb / 5, + sampled: true, + outcome: ResourceOutcome::KilledMemory.into(), + enforced: true, + reserved: Some(Reservation { + memory_kb: reserved_kb, + ..Default::default() + }), + ..Default::default() + }, + ) + .await?; + worker_scheduler + .update_action( + &worker_id, + &operation_id, + UpdateOperationType::UpdateWithActionStage(ActionStage::Completed(ActionResult { + exit_code: 9, + error: Some(make_err!( + Code::FailedPrecondition, + "Action exceeded its memory reservation: reserved {reserved_kb} KiB" + )), + ..ActionResult::default() + })), + ) + .await?; + tokio::task::yield_now().await; + Ok(()) +} + +async fn wait_for_completion(listener: &mut Box) -> ActionResult { + loop { + let (state, _) = listener.changed().await.expect("state stream open"); + if let ActionStage::Completed(result) = &state.stage { + return result.clone(); + } + } +} + +/// A memory kill doubles the reservation and runs the action again; the +/// doubling stops at the largest memory any worker advertises; a kill at +/// that ceiling fails the action with the worker's own error. +#[nativelink_test] +async fn a_memory_kill_requeues_the_action_with_double_the_reservation() -> Result<(), Error> { + let (scheduler, worker_scheduler) = make_scheduler(5); + let mut rx = add_worker(&scheduler, 5_000).await?; + + let base = make_base_action_info( + UNIX_EPOCH + MockClock::time(), + DigestInfo::new([1; 32], 512), + ); + let action_info = Arc::new(ActionInfo { + platform_properties: HashMap::from([("memory_kb".to_string(), "2000".to_string())]), + ..(*base).clone() + }); + let mut listener = scheduler + .add_action(OperationId::default(), action_info) + .await?; + tokio::task::yield_now().await; + + let (first, memory) = next_dispatch(&mut rx).await; + assert_eq!(memory, 2_000); + report_memory_kill(worker_scheduler.as_ref(), &first.operation_id, 2_000).await?; + + let (second, memory) = next_dispatch(&mut rx).await; + assert_eq!(memory, 4_000, "the reservation doubles after a kill"); + assert_eq!( + second.operation_id, first.operation_id, + "same operation, asking for more" + ); + report_memory_kill(worker_scheduler.as_ref(), &second.operation_id, 4_000).await?; + + let (third, memory) = next_dispatch(&mut rx).await; + assert_eq!(memory, 5_000, "capped at the largest worker"); + report_memory_kill(worker_scheduler.as_ref(), &third.operation_id, 5_000).await?; + + // Nothing bigger exists: the client gets the worker's error. + let result = wait_for_completion(&mut listener).await; + let err = result.error.expect("the action fails"); + assert_eq!(err.code, Code::FailedPrecondition, "{err}"); + assert!( + err.to_string().contains("memory reservation"), + "the worker's message reaches the client: {err}" + ); + Ok(()) +} + +/// Escalations have their own budget: four kills on a `max_job_retries` of +/// one, and the action is still running, each time with more memory. +#[nativelink_test] +async fn escalation_does_not_spend_the_retry_budget() -> Result<(), Error> { + // One retry for the action's own failures; escalations are not those. + let (scheduler, worker_scheduler) = make_scheduler(1); + let mut rx = add_worker(&scheduler, 1_000_000).await?; + + let base = make_base_action_info( + UNIX_EPOCH + MockClock::time(), + DigestInfo::new([2; 32], 512), + ); + let action_info = Arc::new(ActionInfo { + platform_properties: HashMap::from([("memory_kb".to_string(), "1000".to_string())]), + ..(*base).clone() + }); + let _listener = scheduler + .add_action(OperationId::default(), action_info) + .await?; + tokio::task::yield_now().await; + + let mut reserved = 1_000; + for _ in 0..4 { + let (dispatch, memory) = next_dispatch(&mut rx).await; + assert_eq!(memory, reserved); + report_memory_kill(worker_scheduler.as_ref(), &dispatch.operation_id, reserved).await?; + reserved *= 2; + } + // Four kills, four requeues, on a budget of one: still running. + let (_, memory) = next_dispatch(&mut rx).await; + assert_eq!(memory, 16_000); + Ok(()) +} + +/// `max_steps` bounds the escalations: with two, the third kill ends the +/// action, and the client hears that the budget is spent. +#[nativelink_test] +async fn escalation_stops_at_max_steps() -> Result<(), Error> { + let (scheduler, worker_scheduler) = make_scheduler_with_limits(5, vec![], 2); + let mut rx = add_worker(&scheduler, 1_000_000).await?; + + let base = make_base_action_info( + UNIX_EPOCH + MockClock::time(), + DigestInfo::new([4; 32], 512), + ); + let action_info = Arc::new(ActionInfo { + platform_properties: HashMap::from([("memory_kb".to_string(), "1000".to_string())]), + ..(*base).clone() + }); + let mut listener = scheduler + .add_action(OperationId::default(), action_info) + .await?; + tokio::task::yield_now().await; + + let mut reserved = 1_000; + for _ in 0..3 { + let (dispatch, memory) = next_dispatch(&mut rx).await; + assert_eq!(memory, reserved); + report_memory_kill(worker_scheduler.as_ref(), &dispatch.operation_id, reserved).await?; + reserved *= 2; + } + + let result = wait_for_completion(&mut listener).await; + let err = result.error.expect("the third kill is past the budget"); + assert_eq!(err.code, Code::FailedPrecondition, "{err}"); + assert!(err.to_string().contains("more than max_steps"), "{err}"); + Ok(()) +} + +/// A worker lost mid-action is not the action's failure: the requeue has +/// its own budget (`max_worker_loss_retries`, two here) and leaves +/// `max_job_retries` alone; past that budget the client hears the worker +/// keeps dying under this action. +#[nativelink_test] +async fn a_lost_worker_does_not_spend_the_retry_budget() -> Result<(), Error> { + let (scheduler, worker_scheduler) = make_scheduler(1); + let mut rx = add_worker(&scheduler, 1_000_000).await?; + + let base = make_base_action_info( + UNIX_EPOCH + MockClock::time(), + DigestInfo::new([4; 32], 512), + ); + let action_info = Arc::new(ActionInfo { + platform_properties: HashMap::from([("memory_kb".to_string(), "1000".to_string())]), + ..(*base).clone() + }); + let mut listener = scheduler + .add_action(OperationId::default(), action_info) + .await?; + tokio::task::yield_now().await; + + for loss in 1..=2 { + let (_, memory) = next_dispatch(&mut rx).await; + assert_eq!(memory, 1_000, "loss {loss}: the reservation is unchanged"); + worker_scheduler + .remove_worker(&WorkerId(WORKER.to_string())) + .await?; + tokio::task::yield_now().await; + rx = add_worker(&scheduler, 1_000_000).await?; + } + // Two losses on a job budget of one: still queued and dispatched. + let (_, memory) = next_dispatch(&mut rx).await; + assert_eq!(memory, 1_000); + // The third loss is past the loss budget. + worker_scheduler + .remove_worker(&WorkerId(WORKER.to_string())) + .await?; + let result = wait_for_completion(&mut listener).await; + let err = result.error.expect("the action fails past the loss budget"); + assert_eq!(err.code, Code::FailedPrecondition, "{err}"); + assert!(err.to_string().contains("lost 3 times"), "{err}"); + Ok(()) +} + +/// With a size ladder, a kill steps to the next class above the reservation +/// rather than scaling; past the top class the largest worker is reserved +/// whole, and only a kill there fails the action. +#[nativelink_test] +async fn escalation_steps_the_ladder_then_takes_the_whole_worker() -> Result<(), Error> { + let (scheduler, worker_scheduler) = make_scheduler_with_ladder(5, vec![1_000, 4_000, 9_000]); + let mut rx = add_worker(&scheduler, 1_000_000).await?; + + let base = make_base_action_info( + UNIX_EPOCH + MockClock::time(), + DigestInfo::new([3; 32], 512), + ); + let action_info = Arc::new(ActionInfo { + platform_properties: HashMap::from([("memory_kb".to_string(), "1500".to_string())]), + ..(*base).clone() + }); + let mut listener = scheduler + .add_action(OperationId::default(), action_info) + .await?; + tokio::task::yield_now().await; + + let (first, memory) = next_dispatch(&mut rx).await; + assert_eq!(memory, 1_500); + report_memory_kill(worker_scheduler.as_ref(), &first.operation_id, 1_500).await?; + let (second, memory) = next_dispatch(&mut rx).await; + assert_eq!(memory, 4_000, "the next class above 1500"); + report_memory_kill(worker_scheduler.as_ref(), &second.operation_id, 4_000).await?; + let (third, memory) = next_dispatch(&mut rx).await; + assert_eq!(memory, 9_000, "the top class"); + report_memory_kill(worker_scheduler.as_ref(), &third.operation_id, 9_000).await?; + let (fourth, memory) = next_dispatch(&mut rx).await; + assert_eq!(memory, 1_000_000, "the whole of the largest worker"); + report_memory_kill(worker_scheduler.as_ref(), &fourth.operation_id, 1_000_000).await?; + + let result = wait_for_completion(&mut listener).await; + let err = result.error.expect("nothing above the whole worker"); + assert_eq!(err.code, Code::FailedPrecondition, "{err}"); + assert!(err.to_string().contains("no worker can run it"), "{err}"); + Ok(()) +} diff --git a/nativelink-scheduler/tests/fleet_capabilities_test.rs b/nativelink-scheduler/tests/fleet_capabilities_test.rs index cc8a773dace..ff2c8c23496 100644 --- a/nativelink-scheduler/tests/fleet_capabilities_test.rs +++ b/nativelink-scheduler/tests/fleet_capabilities_test.rs @@ -166,7 +166,7 @@ async fn scheduler_defers_to_a_peers_worker_it_cannot_see() -> Result<(), Error> let cpu_worker = shape(0); let gpu_worker = shape(1); - let (cpu_tx, _cpu_rx) = mpsc::unbounded_channel(); + let (cpu_tx, _cpu_rx) = mpsc::channel(64); scheduler_a .add_worker(Worker::new( WorkerId("cpu".to_string()), @@ -176,7 +176,7 @@ async fn scheduler_defers_to_a_peers_worker_it_cannot_see() -> Result<(), Error> 1, )) .await?; - let (gpu_tx, _gpu_rx) = mpsc::unbounded_channel(); + let (gpu_tx, _gpu_rx) = mpsc::channel(64); scheduler_b .add_worker(Worker::new( WorkerId("gpu".to_string()), diff --git a/nativelink-scheduler/tests/historical_resource_scheduler_test.rs b/nativelink-scheduler/tests/historical_resource_scheduler_test.rs index 08a92349aef..28e375dacd1 100644 --- a/nativelink-scheduler/tests/historical_resource_scheduler_test.rs +++ b/nativelink-scheduler/tests/historical_resource_scheduler_test.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use core::time::Duration; use std::collections::HashMap; use std::fs; use std::sync::Arc; @@ -22,7 +23,10 @@ mod utils { } use futures::join; -use nativelink_config::schedulers::{HistoricalResourceSpec, SchedulerSpec, SimpleSpec}; +use nativelink_config::schedulers::{ + ColdStartSpec, HistoricalResourceSpec, SchedulerSpec, SimpleSpec, SizeClassSpec, + TimeoutClassSpec, +}; use nativelink_error::Error; use nativelink_macro::nativelink_test; use nativelink_proto::build::bazel::remote::execution::v2::RequestMetadata; @@ -49,18 +53,262 @@ fn write_hints_file(contents: &str) -> String { } fn make_scheduler(hints_file: String) -> (Arc, HistoricalResourceScheduler) { - let mock_scheduler = Arc::new(MockActionScheduler::new()); - let spec = HistoricalResourceSpec { + make_scheduler_with_cold_start(hints_file, None) +} + +fn make_scheduler_with_cold_start( + hints_file: String, + cold_start: Option, +) -> (Arc, HistoricalResourceScheduler) { + make_scheduler_with_spec(&base_spec(hints_file, cold_start)) +} + +fn base_spec(hints_file: String, cold_start: Option) -> HistoricalResourceSpec { + HistoricalResourceSpec { hints_file, refresh_interval_s: 0, cpu_property_name: "cpu_count".to_string(), memory_property_name: "memory_kb".to_string(), + disk_property_name: "disk_kb".to_string(), + classes: vec![], + default_class: None, + class_by_mnemonic: HashMap::new(), + class_by_timeout: vec![], + cold_start, scheduler: Box::new(SchedulerSpec::Simple(SimpleSpec::default())), - }; - let scheduler = HistoricalResourceScheduler::new(&spec, mock_scheduler.clone()); + } +} + +fn make_scheduler_with_spec( + spec: &HistoricalResourceSpec, +) -> (Arc, HistoricalResourceScheduler) { + let mock_scheduler = Arc::new(MockActionScheduler::new()); + let scheduler = HistoricalResourceScheduler::new(spec, mock_scheduler.clone()); (mock_scheduler, scheduler) } +fn class(name: &str, cpu_count: u64, memory_kb: u64, disk_kb: u64) -> SizeClassSpec { + SizeClassSpec { + name: name.to_string(), + cpu_count, + memory_kb, + disk_kb, + } +} + +/// Runs `add_action` under the given Bazel metadata (none when `None`) and +/// returns the platform properties the nested scheduler received. +async fn properties_after_add( + scheduler: &HistoricalResourceScheduler, + mock_scheduler: &MockActionScheduler, + action_info: Arc, + metadata: Option<(&str, &str)>, +) -> Result, Error> { + let client_operation_id = OperationId::default(); + let (_tx, rx) = watch::channel(Arc::new(ActionState { + client_operation_id: OperationId::default(), + stage: ActionStage::Queued, + action_digest: action_info.unique_qualifier.digest(), + last_transition_timestamp: SystemTime::now(), + })); + let context = match metadata { + Some((target_id, action_mnemonic)) => Context::current_with_baggage(vec![KeyValue::new( + BAZEL_METADATA_KEY, + request_metadata_to_baggage(&RequestMetadata { + target_id: target_id.to_string(), + action_mnemonic: action_mnemonic.to_string(), + ..Default::default() + }), + )]), + None => Context::new(), + }; + let (_, (_, passed_action_info)) = join!( + scheduler + .add_action(client_operation_id.clone(), action_info.clone()) + .with_context(context), + mock_scheduler.expect_add_action(Ok(Box::new(TokioWatchActionStateResult::new( + client_operation_id, + action_info, + rx, + )))), + ); + Ok(passed_action_info.platform_properties) +} + +/// A hint keyed by the command digest reaches an action that carries no +/// Bazel metadata at all, in either digest spelling. +#[nativelink_test] +async fn command_digest_hint_matches_an_action_without_metadata() -> Result<(), Error> { + let digest = DigestInfo::new([7; 32], 10); + let hints_file = write_hints_file(&format!( + r#"[ {{ "command_digest": "{}", "cpu_count": 3000, "memory_kb": 8000000, "disk_kb": 500000 }} ]"#, + digest.to_string().replace('-', "/") + )); + let (mock_scheduler, scheduler) = make_scheduler(hints_file.clone()); + let mut action_info = make_base_action_info(UNIX_EPOCH, DigestInfo::zero_digest()) + .as_ref() + .clone(); + action_info.command_digest = digest; + let properties = + properties_after_add(&scheduler, &mock_scheduler, Arc::new(action_info), None).await?; + assert_eq!( + properties, + HashMap::from([ + ("cpu_count".to_string(), "3000".to_string()), + ("memory_kb".to_string(), "8000000".to_string()), + ("disk_kb".to_string(), "500000".to_string()), + ]) + ); + drop(fs::remove_file(hints_file)); + Ok(()) +} + +/// A hint keyed by the action digest, which is what origin events carry, +/// matches the action whose digest it is and no other. +#[nativelink_test] +async fn action_digest_hint_matches_only_its_action() -> Result<(), Error> { + let digest = DigestInfo::new([9; 32], 42); + let hints_file = write_hints_file(&format!( + r#"[ {{ "action_digest": "{digest}", "memory_kb": 7000000 }} ]"# + )); + let (mock_scheduler, scheduler) = make_scheduler(hints_file.clone()); + let matching = make_base_action_info(UNIX_EPOCH, digest); + let properties = properties_after_add(&scheduler, &mock_scheduler, matching, None).await?; + assert_eq!(properties["memory_kb"], "7000000"); + let other = make_base_action_info(UNIX_EPOCH, DigestInfo::new([8; 32], 42)); + let properties = properties_after_add(&scheduler, &mock_scheduler, other, None).await?; + assert!(properties.is_empty(), "{properties:?}"); + drop(fs::remove_file(hints_file)); + Ok(()) +} + +/// A hint may name a class instead of numbers; an untagged action takes the +/// class the cold-start policy picks: by mnemonic, else by timeout, else the +/// default. A value the client sent is never lowered. +#[nativelink_test] +async fn classes_resolve_hints_and_cold_start() -> Result<(), Error> { + let hints_file = write_hints_file( + r#"[ { "action_mnemonic": "TestRunner", "class": "m", "cpu_count": 5000 } ]"#, + ); + let mut spec = base_spec(hints_file.clone(), None); + spec.classes = vec![ + class("s", 1000, 4_000_000, 1_000_000), + class("m", 2000, 12_000_000, 4_000_000), + class("l", 4000, 36_000_000, 16_000_000), + ]; + spec.default_class = Some("s".to_string()); + spec.class_by_mnemonic = HashMap::from([("Link".to_string(), "l".to_string())]); + spec.class_by_timeout = vec![TimeoutClassSpec { + min_timeout_s: 900, + class: "m".to_string(), + }]; + let (mock_scheduler, scheduler) = make_scheduler_with_spec(&spec); + let base = make_base_action_info(UNIX_EPOCH, DigestInfo::zero_digest()); + + // A hint with a class: the class fills what the hint leaves out. + let properties = properties_after_add( + &scheduler, + &mock_scheduler, + base.clone(), + Some(("//pkg:t", "TestRunner")), + ) + .await?; + assert_eq!( + properties["cpu_count"], "5000", + "the hint's own number wins" + ); + assert_eq!( + properties["memory_kb"], "12000000", + "the class fills memory" + ); + assert_eq!(properties["disk_kb"], "4000000"); + + // No hint, a known heavy mnemonic: its class. + let properties = properties_after_add( + &scheduler, + &mock_scheduler, + base.clone(), + Some(("//pkg:x", "Link")), + ) + .await?; + assert_eq!(properties["memory_kb"], "36000000"); + + // No hint, an unknown mnemonic, a long timeout: the timeout rule. + let mut long = (*base).clone(); + long.timeout = Duration::from_hours(1); + let properties = properties_after_add( + &scheduler, + &mock_scheduler, + Arc::new(long), + Some(("//pkg:y", "Other")), + ) + .await?; + assert_eq!(properties["memory_kb"], "12000000"); + + // Nothing at all: the default class, but a client value stands. + let mut tagged = (*base).clone(); + // The helper's timeout is MAX, which every timeout rule reaches; give it + // one none does, so the default class is what gets tested. + tagged.timeout = Duration::from_secs(10); + tagged + .platform_properties + .insert("memory_kb".to_string(), "9".to_string()); + let properties = + properties_after_add(&scheduler, &mock_scheduler, Arc::new(tagged), None).await?; + assert_eq!(properties["memory_kb"], "9"); + assert_eq!(properties["cpu_count"], "1000"); + assert_eq!(properties["disk_kb"], "1000000"); + drop(fs::remove_file(hints_file)); + Ok(()) +} + +/// An untagged action nothing hints at gets the cold-start reservation; +/// a value the client sent is left as it is. +#[nativelink_test] +async fn cold_start_fills_only_absent_properties() -> Result<(), Error> { + let hints_file = write_hints_file("[]"); + let (mock_scheduler, scheduler) = make_scheduler_with_cold_start( + hints_file.clone(), + Some(ColdStartSpec { + cpu_count: 1000, + memory_kb: 4_000_000, + disk_kb: 0, + }), + ); + let client_operation_id = OperationId::default(); + let mut action_info = make_base_action_info(UNIX_EPOCH, DigestInfo::zero_digest()); + Arc::make_mut(&mut action_info) + .platform_properties + .insert("memory_kb".to_string(), "1000".to_string()); + let (forward_watch_channel_tx, forward_watch_channel_rx) = + watch::channel(Arc::new(ActionState { + client_operation_id: client_operation_id.clone(), + stage: ActionStage::Queued, + action_digest: action_info.unique_qualifier.digest(), + last_transition_timestamp: SystemTime::UNIX_EPOCH, + })); + drop(forward_watch_channel_tx); + + let (_, (_, passed_action_info)) = join!( + scheduler.add_action(client_operation_id.clone(), action_info.clone()), + mock_scheduler.expect_add_action(Ok(Box::new(TokioWatchActionStateResult::new( + client_operation_id.clone(), + action_info, + forward_watch_channel_rx, + )))), + ); + + assert_eq!( + HashMap::from([ + ("cpu_count".to_string(), "1000".to_string()), + ("memory_kb".to_string(), "1000".to_string()), + ]), + passed_action_info.platform_properties + ); + drop(fs::remove_file(hints_file)); + Ok(()) +} + #[nativelink_test] async fn add_action_applies_historical_resource_hint() -> Result<(), Error> { let hints_file = write_hints_file( @@ -126,3 +374,101 @@ async fn add_action_applies_historical_resource_hint() -> Result<(), Error> { drop(fs::remove_file(hints_file)); Ok(()) } + +/// A record naming a target and a digest is reachable by either key; a +/// single composite key was reachable by neither. +#[nativelink_test] +async fn a_hint_naming_a_target_and_a_digest_is_found_by_either() -> Result<(), Error> { + let digest = DigestInfo::new([3; 32], 12); + let hints_file = write_hints_file(&format!( + r#"[ {{ "target_id": "//pkg:t", "action_digest": "{digest}", "memory_kb": 5000000 }} ]"# + )); + let (mock_scheduler, scheduler) = make_scheduler(hints_file.clone()); + let by_digest = make_base_action_info(UNIX_EPOCH, digest); + let properties = properties_after_add(&scheduler, &mock_scheduler, by_digest, None).await?; + assert_eq!(properties["memory_kb"], "5000000", "found by the digest"); + let by_target = make_base_action_info(UNIX_EPOCH, DigestInfo::new([4; 32], 12)); + let properties = properties_after_add( + &scheduler, + &mock_scheduler, + by_target, + Some(("//pkg:t", "Other")), + ) + .await?; + assert_eq!(properties["memory_kb"], "5000000", "found by the target"); + drop(fs::remove_file(hints_file)); + Ok(()) +} + +/// The ladder holds: an action with a mnemonic and no target takes the +/// digest hint before the mnemonic one. +#[nativelink_test] +async fn a_mnemonic_alone_does_not_outrank_a_digest_hint() -> Result<(), Error> { + let digest = DigestInfo::new([5; 32], 12); + let hints_file = write_hints_file(&format!( + r#"[ {{ "action_mnemonic": "Link", "memory_kb": 1000000 }}, + {{ "action_digest": "{digest}", "memory_kb": 9000000 }} ]"# + )); + let (mock_scheduler, scheduler) = make_scheduler(hints_file.clone()); + let action = make_base_action_info(UNIX_EPOCH, digest); + let properties = + properties_after_add(&scheduler, &mock_scheduler, action, Some(("", "Link"))).await?; + assert_eq!(properties["memory_kb"], "9000000"); + drop(fs::remove_file(hints_file)); + Ok(()) +} + +/// A hint that reserves nothing is no hint: the cold start applies. +#[nativelink_test] +async fn an_all_zero_hint_falls_through_to_the_cold_start() -> Result<(), Error> { + let digest = DigestInfo::new([6; 32], 12); + let hints_file = write_hints_file(&format!( + r#"[ {{ "action_digest": "{digest}", "samples": 3, "last_seen_s": 1 }} ]"# + )); + let (mock_scheduler, scheduler) = make_scheduler_with_cold_start( + hints_file.clone(), + Some(ColdStartSpec { + cpu_count: 1500, + memory_kb: 0, + disk_kb: 250_000, + }), + ); + let action = make_base_action_info(UNIX_EPOCH, digest); + let properties = properties_after_add(&scheduler, &mock_scheduler, action, None).await?; + assert_eq!(properties["cpu_count"], "1500"); + assert_eq!( + properties["disk_kb"], "250000", + "the cold start reserves disk too" + ); + assert!(!properties.contains_key("memory_kb")); + drop(fs::remove_file(hints_file)); + Ok(()) +} + +/// A rule naming a class the ladder lacks is refused at load; at run time +/// the next candidate that exists is taken, never the raw numbers. +#[nativelink_test] +async fn a_rule_naming_no_class_is_refused_and_the_default_class_stands_in() -> Result<(), Error> { + let hints_file = write_hints_file("[]"); + let mut spec = base_spec(hints_file.clone(), None); + spec.classes = vec![class("s", 1000, 4_000_000, 0)]; + spec.default_class = Some("s".to_string()); + spec.class_by_mnemonic = HashMap::from([("Link".to_string(), "xl".to_string())]); + let err = HistoricalResourceScheduler::validate(&spec).expect_err("xl is not on the ladder"); + assert!(err.to_string().contains("class_by_mnemonic[Link]"), "{err}"); + let (mock_scheduler, scheduler) = make_scheduler_with_spec(&spec); + let action = make_base_action_info(UNIX_EPOCH, DigestInfo::zero_digest()); + let properties = properties_after_add( + &scheduler, + &mock_scheduler, + action, + Some(("//pkg:x", "Link")), + ) + .await?; + assert_eq!( + properties["memory_kb"], "4000000", + "the default class, not nothing" + ); + drop(fs::remove_file(hints_file)); + Ok(()) +} diff --git a/nativelink-scheduler/tests/live_memory_veto_test.rs b/nativelink-scheduler/tests/live_memory_veto_test.rs index 374103b8232..30c62e0195b 100644 --- a/nativelink-scheduler/tests/live_memory_veto_test.rs +++ b/nativelink-scheduler/tests/live_memory_veto_test.rs @@ -47,8 +47,8 @@ async fn add_worker( scheduler: &SimpleScheduler, name: &str, properties: HashMap, -) -> Result, Error> { - let (tx, mut rx) = mpsc::unbounded_channel(); +) -> Result, Error> { + let (tx, mut rx) = mpsc::channel(64); let worker = Worker::new( WorkerId(name.to_string()), PlatformProperties::new(properties), @@ -87,7 +87,7 @@ async fn add_action( Ok(()) } -fn dispatched(rx: &mut mpsc::UnboundedReceiver) -> usize { +fn dispatched(rx: &mut mpsc::Receiver) -> usize { let mut count = 0; while let Ok(msg) = rx.try_recv() { if matches!(msg.update, Some(update_for_worker::Update::StartAction(_))) { @@ -175,6 +175,36 @@ async fn live_memory_veto_skips_a_worker_that_reports_no_room() -> Result<(), Er /// A worker that reported no room and then reports room again has to wake /// the matcher itself. With the fallback pass off nothing else would, and /// the vetoed action would wait for an unrelated event. +/// The connection result names the property the veto reads, so the worker +/// declines for load against the same one the scheduler would have vetoed. +#[nativelink_test] +async fn the_connection_result_names_the_veto_property() -> Result<(), Error> { + let scheduler = make_scheduler(&SimpleSpec { + supported_platform_properties: Some(HashMap::from([( + "memory_kb".to_string(), + PropertyType::Minimum, + )])), + live_memory_veto: Some("memory_kb".to_string()), + ..SimpleSpec::default() + }); + let (tx, mut rx) = mpsc::channel(64); + scheduler + .add_worker(Worker::new( + WorkerId("told".to_string()), + PlatformProperties::new(memory_worker(100_000)), + tx, + NOW_TIME, + /* max_inflight_tasks */ 4, + )) + .await?; + let connected = rx.recv().await.unwrap(); + let Some(update_for_worker::Update::ConnectionResult(result)) = connected.update else { + panic!("expected a ConnectionResult, got {connected:?}"); + }; + assert_eq!(result.memory_property, "memory_kb"); + Ok(()) +} + #[nativelink_test] async fn a_load_report_with_room_again_wakes_the_matcher() -> Result<(), Error> { let scheduler = make_scheduler(&SimpleSpec { diff --git a/nativelink-scheduler/tests/placement_test.rs b/nativelink-scheduler/tests/placement_test.rs new file mode 100644 index 00000000000..8a4a6675461 --- /dev/null +++ b/nativelink-scheduler/tests/placement_test.rs @@ -0,0 +1,245 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use core::time::Duration; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::UNIX_EPOCH; + +use mock_instant::thread_local::MockClock; +use nativelink_config::schedulers::{PropertyType, SimpleSpec, WorkerAllocationStrategy}; +use nativelink_error::{Error, ResultExt}; +use nativelink_macro::nativelink_test; +use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{ + UpdateForWorker, update_for_worker, +}; +use nativelink_scheduler::default_scheduler_factory::memory_awaited_action_db_factory; +use nativelink_scheduler::simple_scheduler::SimpleScheduler; +use nativelink_scheduler::worker::Worker; +use nativelink_scheduler::worker_scheduler::WorkerScheduler; +use nativelink_util::action_messages::{ActionInfo, OperationId, WorkerId}; +use nativelink_util::common::DigestInfo; +use nativelink_util::instant_wrapper::MockInstantWrapped; +use nativelink_util::operation_state_manager::ClientStateManager; +use nativelink_util::platform_properties::{PlatformProperties, PlatformPropertyValue}; +use pretty_assertions::assert_eq; +use tokio::sync::{Notify, mpsc}; +use utils::scheduler_utils::make_base_action_info; + +mod utils { + pub(crate) mod scheduler_utils; +} + +const NOW_TIME: u64 = 10000; + +fn make_scheduler(strategy: WorkerAllocationStrategy) -> Arc { + make_scheduler_with_spec(&SimpleSpec { + allocation_strategy: strategy, + ..SimpleSpec::default() + }) +} + +fn make_scheduler_with_spec(spec: &SimpleSpec) -> Arc { + MockClock::set_time(Duration::from_secs(NOW_TIME)); + let task_change_notify = Arc::new(Notify::new()); + let (scheduler, _worker_scheduler) = SimpleScheduler::new_with_callback( + spec, + memory_awaited_action_db_factory(0, &task_change_notify, MockInstantWrapped::default), + || async move {}, + task_change_notify, + MockInstantWrapped::default, + None, + ); + scheduler +} + +async fn add_worker( + scheduler: &SimpleScheduler, + name: &str, +) -> Result, Error> { + add_worker_with_properties(scheduler, name, HashMap::new()).await +} + +async fn add_worker_with_properties( + scheduler: &SimpleScheduler, + name: &str, + properties: HashMap, +) -> Result, Error> { + let (tx, mut rx) = mpsc::channel(64); + let worker = Worker::new( + WorkerId(name.to_string()), + PlatformProperties::new(properties), + tx, + NOW_TIME, + /* max_inflight_tasks */ 4, + ); + scheduler + .add_worker(worker) + .await + .err_tip(|| "Failed to add worker")?; + tokio::task::yield_now().await; + let connected = rx.recv().await.unwrap(); + assert!( + matches!( + connected.update, + Some(update_for_worker::Update::ConnectionResult(_)) + ), + "first message should be the connection result: {connected:?}" + ); + Ok(rx) +} + +async fn add_action(scheduler: &SimpleScheduler, digest_byte: u8) -> Result<(), Error> { + add_action_with_properties(scheduler, digest_byte, HashMap::new()).await +} + +async fn add_action_with_properties( + scheduler: &SimpleScheduler, + digest_byte: u8, + platform_properties: HashMap, +) -> Result<(), Error> { + let base = make_base_action_info( + UNIX_EPOCH + MockClock::time(), + DigestInfo::new([digest_byte; 32], 512), + ); + let action_info = Arc::new(ActionInfo { + platform_properties, + ..(*base).clone() + }); + scheduler + .add_action(OperationId::default(), action_info) + .await?; + tokio::task::yield_now().await; + Ok(()) +} + +fn dispatched(rx: &mut mpsc::Receiver) -> usize { + let mut count = 0; + while let Ok(msg) = rx.try_recv() { + if matches!(msg.update, Some(update_for_worker::Update::StartAction(_))) { + count += 1; + } + } + count +} + +/// One worker takes the first actions; a second worker joins; the next action +/// goes to the idle newcomer under `least_loaded`, and to the busy veteran +/// under the LRU default, which only knows who was used least recently. +async fn run(strategy: WorkerAllocationStrategy) -> Result<(usize, usize), Error> { + let scheduler = make_scheduler(strategy); + let mut veteran = add_worker(&scheduler, "veteran").await?; + add_action(&scheduler, 1).await?; + add_action(&scheduler, 2).await?; + assert_eq!( + dispatched(&mut veteran), + 2, + "the only worker takes everything" + ); + + let mut newcomer = add_worker(&scheduler, "newcomer").await?; + add_action(&scheduler, 3).await?; + Ok((dispatched(&mut veteran), dispatched(&mut newcomer))) +} + +#[nativelink_test] +async fn least_loaded_gives_the_next_action_to_the_idle_newcomer() -> Result<(), Error> { + assert_eq!(run(WorkerAllocationStrategy::LeastLoaded).await?, (0, 1)); + Ok(()) +} + +#[nativelink_test] +async fn least_recently_used_gives_it_to_the_busy_veteran() -> Result<(), Error> { + assert_eq!( + run(WorkerAllocationStrategy::LeastRecentlyUsed).await?, + (1, 0) + ); + Ok(()) +} + +fn memory_worker(kb: u64) -> HashMap { + HashMap::from([("memory_kb".to_string(), PlatformPropertyValue::Minimum(kb))]) +} + +fn memory_action(kb: u64) -> HashMap { + HashMap::from([("memory_kb".to_string(), kb.to_string())]) +} + +fn worker_props(memory_kb: u64, zone: Option<&str>) -> HashMap { + let mut props = memory_worker(memory_kb); + if let Some(zone) = zone { + props.insert( + "zone".to_string(), + PlatformPropertyValue::Priority(zone.to_string()), + ); + } + props +} + +fn best_fit_scheduler() -> Arc { + make_scheduler_with_spec(&SimpleSpec { + supported_platform_properties: Some(HashMap::from([ + ("memory_kb".to_string(), PropertyType::Minimum), + ("zone".to_string(), PropertyType::Priority), + ])), + allocation_strategy: WorkerAllocationStrategy::BestFit, + ..SimpleSpec::default() + }) +} + +/// A 4 GiB action goes to the worker it fills best: the one with 6 GiB +/// left, not the one with 30 GiB left, whatever the order they joined in. +#[nativelink_test] +async fn best_fit_keeps_the_big_room_for_big_actions() -> Result<(), Error> { + let scheduler = best_fit_scheduler(); + let mut roomy = + add_worker_with_properties(&scheduler, "roomy", worker_props(30_000, None)).await?; + let mut snug = + add_worker_with_properties(&scheduler, "snug", worker_props(6_000, None)).await?; + add_action_with_properties(&scheduler, 1, memory_action(4_000)).await?; + assert_eq!((dispatched(&mut roomy), dispatched(&mut snug)), (0, 1)); + + // The next 4 GiB action no longer fits the snug worker (2 GiB left) and + // takes the roomy one. + add_action_with_properties(&scheduler, 2, memory_action(4_000)).await?; + assert_eq!((dispatched(&mut roomy), dispatched(&mut snug)), (1, 0)); + Ok(()) +} + +/// A priority property is a preference: the worker carrying the action's +/// zone is chosen first, even when another fits more tightly; an action +/// naming a zone no worker has still runs, on the tightest fit. +#[nativelink_test] +async fn best_fit_prefers_the_worker_with_the_actions_priority_value() -> Result<(), Error> { + let scheduler = best_fit_scheduler(); + let mut east = + add_worker_with_properties(&scheduler, "east", worker_props(30_000, Some("east"))).await?; + let mut west = + add_worker_with_properties(&scheduler, "west", worker_props(6_000, Some("west"))).await?; + + let mut wants_east = memory_action(4_000); + wants_east.insert("zone".to_string(), "east".to_string()); + add_action_with_properties(&scheduler, 3, wants_east).await?; + assert_eq!((dispatched(&mut east), dispatched(&mut west)), (1, 0)); + + let mut wants_north = memory_action(4_000); + wants_north.insert("zone".to_string(), "north".to_string()); + add_action_with_properties(&scheduler, 4, wants_north).await?; + assert_eq!( + (dispatched(&mut east), dispatched(&mut west)), + (0, 1), + "no worker is north; the tightest fit takes it" + ); + Ok(()) +} diff --git a/nativelink-scheduler/tests/redis_store_awaited_action_db_test.rs b/nativelink-scheduler/tests/redis_store_awaited_action_db_test.rs index 460ee445899..c13978bb319 100644 --- a/nativelink-scheduler/tests/redis_store_awaited_action_db_test.rs +++ b/nativelink-scheduler/tests/redis_store_awaited_action_db_test.rs @@ -33,6 +33,7 @@ use nativelink_redis_tester::FakeRedisBackend; use nativelink_scheduler::awaited_action_db::{ AwaitedAction, AwaitedActionDb, AwaitedActionSubscriber, }; +use nativelink_scheduler::default_scheduler_factory::retain_completed_for_s; use nativelink_scheduler::simple_scheduler::SimpleScheduler; use nativelink_scheduler::store_awaited_action_db::StoreAwaitedActionDb; use nativelink_scheduler::worker::Worker; @@ -50,7 +51,6 @@ use nativelink_util::store_trait::SchedulerStore; use parking_lot::Mutex; use pretty_assertions::assert_eq; use redis::Value; -use tokio::sync::mpsc::unbounded_channel; use tokio::sync::{Notify, mpsc}; use tonic::Code; use utils::scheduler_utils::update_eq; @@ -63,13 +63,15 @@ const INSTANCE_NAME: &str = "instance_name"; async fn verify_initial_connection_message( worker_id: WorkerId, - rx: &mut mpsc::UnboundedReceiver, + rx: &mut mpsc::Receiver, ) { // Worker should have been sent an execute command. let expected_msg_for_worker = UpdateForWorker { update: Some(update_for_worker::Update::ConnectionResult( ConnectionResult { worker_id: worker_id.into(), + dispatch_ack: true, + memory_property: String::new(), }, )), }; @@ -83,8 +85,8 @@ async fn setup_new_worker( scheduler: &SimpleScheduler, worker_id: WorkerId, props: PlatformProperties, -) -> Result, Error> { - let (tx, mut rx) = unbounded_channel(); +) -> Result, Error> { + let (tx, mut rx) = mpsc::channel(64); let worker = Worker::new(worker_id.clone(), props, tx, NOW_TIME, 0); scheduler .add_worker(worker) @@ -712,3 +714,12 @@ async fn an_update_against_a_vanished_record_leaves_nothing_behind() -> Result<( Ok(()) } + +/// The Redis backend gets the same retention default as the memory backend. +/// Passed through as 0 it wrote every completed record without an expiry. +#[nativelink_test] +async fn unset_retention_is_the_shared_default_not_forever() -> Result<(), Error> { + assert_eq!(retain_completed_for_s(0), 60); + assert_eq!(retain_completed_for_s(30), 30); + Ok(()) +} diff --git a/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs b/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs index 7e6a153751b..eb96633b6cd 100644 --- a/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs +++ b/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs @@ -7,14 +7,15 @@ use std::time::SystemTime; use futures::{Stream, StreamExt, stream}; use mock_instant::thread_local::MockClock; use nativelink_config::schedulers::SimpleSpec; -use nativelink_error::{Code, Error, make_err}; +use nativelink_error::{Code, Error, ResultExt, make_err}; use nativelink_macro::nativelink_test; use nativelink_metric::{ MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent, }; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::UpdateForWorker; use nativelink_scheduler::awaited_action_db::{ - AwaitedAction, AwaitedActionDb, SortedAwaitedAction, SortedAwaitedActionState, + AwaitedAction, AwaitedActionDb, AwaitedActionSubscriber, SortedAwaitedAction, + SortedAwaitedActionState, }; use nativelink_scheduler::default_scheduler_factory::memory_awaited_action_db_factory; use nativelink_scheduler::simple_scheduler::SimpleScheduler; @@ -469,12 +470,44 @@ async fn the_sweep_retires_queued_actions_no_client_is_waiting_on() -> Result<() /// What a listing of the queue does after its first entry. #[derive(Clone, Copy)] enum Fault { - /// The second entry's record is gone, as after an eviction: the listing - /// names it and the read fails. The rest of the listing follows. + /// The listing names a second entry whose record is gone, as after an + /// eviction: the stream yields the read error in its place. The rest of + /// the listing follows. LostRecord, /// The next page fails to load, as when the store goes away mid-read: /// the stream yields the error and ends. Transport, + /// The second entry is listed but reading its record fails with this + /// code, as when it expired or was evicted between the search and the + /// read (`NotFound`), was damaged (`InvalidArgument`), or the store + /// went away (`Unavailable`). + BrokenRecord(Code), + /// The store hands back a row it could not decode, as the Redis listing + /// does for damaged JSON: an `InvalidArgument` item between two good + /// ones, before any subscriber exists. + UndecodableRow, +} + +/// A subscriber whose record may be unreadable. +enum MaybeBroken { + Sound(S), + Broken(Code), +} + +impl AwaitedActionSubscriber for MaybeBroken { + async fn changed(&mut self) -> Result { + match self { + Self::Sound(subscriber) => subscriber.changed().await, + Self::Broken(code) => Err(make_err!(*code, "record unreadable")), + } + } + + async fn borrow(&self) -> Result { + match self { + Self::Sound(subscriber) => subscriber.borrow().await, + Self::Broken(code) => Err(make_err!(*code, "record unreadable")), + } + } } /// The memory database with one fault injected into every queue listing. @@ -493,29 +526,72 @@ impl MetricsComponent for FaultyDb { } } +impl FaultyDb { + /// The listing as the store hands it back, with the fault in place of + /// or beside its second entry. + fn inject( + &self, + listed: Vec>, + ) -> Vec, Error>> { + let mut items = Vec::new(); + for (i, item) in listed.into_iter().enumerate() { + if i == 1 { + match self.fault { + Fault::LostRecord => { + items.push(Err(make_err!(Code::NotFound, "record evicted"))); + } + Fault::UndecodableRow => { + items.push(Err(make_err!( + Code::InvalidArgument, + "In AwaitedAction::decode" + ))); + } + Fault::Transport => { + items.push(Err(make_err!(Code::Unavailable, "FT.CURSOR READ failed"))); + break; + } + Fault::BrokenRecord(code) => { + items.push(Ok(MaybeBroken::Broken(code))); + continue; + } + } + } + items.push(item.map(MaybeBroken::Sound)); + } + items + } +} + impl AwaitedActionDb for FaultyDb { - type Subscriber = T::Subscriber; + type Subscriber = MaybeBroken; async fn get_awaited_action_by_id( &self, client_operation_id: &OperationId, ) -> Result, Error> { - self.inner + Ok(self + .inner .get_awaited_action_by_id(client_operation_id) - .await + .await? + .map(MaybeBroken::Sound)) } async fn get_all_awaited_actions( &self, ) -> Result> + Send, Error> { - self.inner.get_all_awaited_actions().await + let listed: Vec<_> = self.inner.get_all_awaited_actions().await?.collect().await; + Ok(stream::iter(self.inject(listed))) } async fn get_by_operation_id( &self, operation_id: &OperationId, ) -> Result, Error> { - self.inner.get_by_operation_id(operation_id).await + Ok(self + .inner + .get_by_operation_id(operation_id) + .await? + .map(MaybeBroken::Sound)) } async fn get_range_of_actions( @@ -531,22 +607,7 @@ impl AwaitedActionDb for FaultyDb { .await? .collect() .await; - let mut items = Vec::new(); - for (i, item) in listed.into_iter().enumerate() { - if i == 1 { - match self.fault { - Fault::LostRecord => { - items.push(Err(make_err!(Code::NotFound, "record evicted"))); - } - Fault::Transport => { - items.push(Err(make_err!(Code::Unavailable, "FT.CURSOR READ failed"))); - break; - } - } - } - items.push(item); - } - Ok(stream::iter(items)) + Ok(stream::iter(self.inject(listed))) } async fn update_awaited_action(&self, new_awaited_action: AwaitedAction) -> Result<(), Error> { @@ -559,9 +620,11 @@ impl AwaitedActionDb for FaultyDb { action_info: Arc, no_event_action_timeout: Duration, ) -> Result { - self.inner - .add_action(client_operation_id, action_info, no_event_action_timeout) - .await + Ok(MaybeBroken::Sound( + self.inner + .add_action(client_operation_id, action_info, no_event_action_timeout) + .await?, + )) } } @@ -643,6 +706,177 @@ async fn the_sweep_fails_when_the_listing_fails() -> Result<(), Error> { Ok(()) } +/// Queues `n` actions and returns their client listeners. +async fn add_queued_actions( + state_mgr: &impl ClientStateManager, + n: usize, +) -> Result>, Error> { + let mut clients = Vec::with_capacity(n); + for _ in 0..n { + clients.push( + state_mgr + .add_action( + OperationId::default(), + Arc::new(action_info(make_system_time(0))), + ) + .await?, + ); + } + Ok(clients) +} + +/// Three abandoned actions, the second one's record unreadable when the +/// sweep goes to read it. The sweep retires the other two. +#[nativelink_test] +async fn the_sweep_skips_a_record_it_cannot_read() -> Result<(), Error> { + MockClock::set_time(Duration::from_secs(NOW_TIME)); + let state_mgr = faulty_state_manager(Fault::BrokenRecord(Code::NotFound)); + for _ in 0..3 { + drop( + state_mgr + .add_action( + OperationId::default(), + Arc::new(action_info(make_system_time(0))), + ) + .await?, + ); + } + MockClock::advance(Duration::from_mins(6)); + + assert_eq!(state_mgr.sweep_abandoned_queued_actions().await?, 2); + Ok(()) +} + +/// A listing that meets a record it cannot decode leaves it out and goes on, +/// so the provisioner and the admin API, which read this listing every few +/// seconds, still see every other queued action. +#[nativelink_test] +async fn the_listing_skips_a_record_it_cannot_read() -> Result<(), Error> { + MockClock::set_time(Duration::from_secs(NOW_TIME)); + let state_mgr = faulty_state_manager(Fault::BrokenRecord(Code::InvalidArgument)); + // Kept so the clients count as waiting. + let _clients = add_queued_actions(state_mgr.as_ref(), 3).await?; + + let listed = MatchingEngineStateManager::filter_operations( + state_mgr.as_ref(), + OperationFilter { + stages: OperationStageFlags::Queued, + ..Default::default() + }, + ) + .await? + .collect::>() + .await; + assert_eq!(listed.len(), 2, "the unreadable record is left out"); + for operation in &listed { + operation + .as_state() + .await + .err_tip(|| "every listed operation is readable")?; + } + Ok(()) +} + +/// The Redis listing reports a row it cannot decode as an error item, before +/// any record is read. The queued listing leaves it out and lists the rest. +#[nativelink_test] +async fn the_listing_skips_a_row_it_cannot_decode() -> Result<(), Error> { + MockClock::set_time(Duration::from_secs(NOW_TIME)); + let state_mgr = faulty_state_manager(Fault::UndecodableRow); + let _clients = add_queued_actions(state_mgr.as_ref(), 3).await?; + + let listed = MatchingEngineStateManager::filter_operations( + state_mgr.as_ref(), + OperationFilter { + stages: OperationStageFlags::Queued, + ..Default::default() + }, + ) + .await? + .collect::>() + .await; + assert_eq!( + listed.len(), + 3, + "the undecodable row is left out, the rest listed" + ); + for operation in &listed { + operation + .as_state() + .await + .err_tip(|| "every listed operation is readable")?; + } + Ok(()) +} + +/// The listing of every operation, which collects the whole set before it +/// filters, also survives a row it cannot decode instead of failing whole. +#[nativelink_test] +async fn the_listing_of_every_operation_skips_a_row_it_cannot_decode() -> Result<(), Error> { + MockClock::set_time(Duration::from_secs(NOW_TIME)); + let state_mgr = faulty_state_manager(Fault::UndecodableRow); + let _clients = add_queued_actions(state_mgr.as_ref(), 3).await?; + + let listed = MatchingEngineStateManager::filter_operations( + state_mgr.as_ref(), + OperationFilter::default(), + ) + .await? + .collect::>() + .await; + assert_eq!(listed.len(), 3); + Ok(()) +} + +/// The listing of every operation still fails when the store does, rather +/// than quietly listing whatever came before the failure. +#[nativelink_test] +async fn the_listing_of_every_operation_fails_when_the_store_does() -> Result<(), Error> { + MockClock::set_time(Duration::from_secs(NOW_TIME)); + let state_mgr = faulty_state_manager(Fault::Transport); + let _clients = add_queued_actions(state_mgr.as_ref(), 3).await?; + + let err = MatchingEngineStateManager::filter_operations( + state_mgr.as_ref(), + OperationFilter::default(), + ) + .await + .err() + .expect("a store failure fails the listing"); + assert_eq!(err.code, Code::Unavailable); + Ok(()) +} + +/// A read that fails because the store is unreachable is not a lost record; +/// the listing reports it instead of quietly leaving the action out. +#[nativelink_test] +async fn the_listing_reports_a_store_failure_on_read() -> Result<(), Error> { + MockClock::set_time(Duration::from_secs(NOW_TIME)); + let state_mgr = faulty_state_manager(Fault::BrokenRecord(Code::Unavailable)); + // Kept so the clients count as waiting. + let _clients = add_queued_actions(state_mgr.as_ref(), 3).await?; + + let listed = MatchingEngineStateManager::filter_operations( + state_mgr.as_ref(), + OperationFilter { + stages: OperationStageFlags::Queued, + ..Default::default() + }, + ) + .await? + .collect::>() + .await; + let mut failures = 0; + for operation in &listed { + if let Err(err) = operation.as_state().await { + assert_eq!(err.code, Code::Unavailable); + failures += 1; + } + } + assert_eq!(failures, 1, "the store failure reaches the reader"); + Ok(()) +} + fn faulty_scheduler(fault: Fault) -> Arc { MockClock::set_time(Duration::from_secs(NOW_TIME)); let task_change_notify = Arc::new(Notify::new()); @@ -670,7 +904,7 @@ fn faulty_scheduler(fault: Fault) -> Arc { async fn three_queued_actions_and_a_worker( scheduler: &SimpleScheduler, -) -> Result, Error> { +) -> Result, Error> { for _ in 0..3 { scheduler .add_action( @@ -679,7 +913,7 @@ async fn three_queued_actions_and_a_worker( ) .await?; } - let (tx, mut rx) = mpsc::unbounded_channel(); + let (tx, mut rx) = mpsc::channel(64); scheduler .add_worker(Worker::new( WorkerId("worker".to_string()), diff --git a/nativelink-scheduler/tests/simple_scheduler_test.rs b/nativelink-scheduler/tests/simple_scheduler_test.rs index a1608c55729..351474edae6 100644 --- a/nativelink-scheduler/tests/simple_scheduler_test.rs +++ b/nativelink-scheduler/tests/simple_scheduler_test.rs @@ -37,7 +37,7 @@ use nativelink_proto::com::github::trace_machina::nativelink::events::{ }; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{ ActionResourceUsage, ConnectionResult, KillOperationRequest, StartExecute, UpdateForWorker, - update_for_worker, + WorkerLoad, update_for_worker, }; use nativelink_scheduler::awaited_action_db::{ AwaitedAction, AwaitedActionDb, AwaitedActionSubscriber, SortedAwaitedAction, @@ -73,18 +73,21 @@ mod utils { async fn verify_initial_connection_message( worker_id: WorkerId, - rx: &mut mpsc::UnboundedReceiver, + rx: &mut mpsc::Receiver, ) { - // Worker should have been sent an execute command. - let expected_msg_for_worker = UpdateForWorker { - update: Some(update_for_worker::Update::ConnectionResult( - ConnectionResult { - worker_id: worker_id.into(), - }, - )), - }; + // The first message names the worker and announces the acknowledgement; + // the memory property follows each test's own veto setting. let msg_for_worker = rx.recv().await.unwrap(); - assert_eq!(msg_for_worker, expected_msg_for_worker); + let Some(update_for_worker::Update::ConnectionResult(ConnectionResult { + worker_id: sent_worker_id, + dispatch_ack, + .. + })) = msg_for_worker.update + else { + panic!("expected a ConnectionResult, got {msg_for_worker:?}"); + }; + assert_eq!(sent_worker_id, String::from(worker_id)); + assert!(dispatch_ack); } const NOW_TIME: u64 = 10000; @@ -99,8 +102,8 @@ async fn setup_new_worker( scheduler: &SimpleScheduler, worker_id: WorkerId, props: PlatformProperties, -) -> Result, Error> { - let (tx, mut rx) = mpsc::unbounded_channel(); +) -> Result, Error> { + let (tx, mut rx) = mpsc::channel(64); let worker = Worker::new(worker_id.clone(), props, tx, NOW_TIME, 0); scheduler .add_worker(worker) @@ -2446,13 +2449,14 @@ async fn worker_retries_on_internal_error_and_fails_test() -> Result<(), Error> /// hiding the cluster-side root cause behind a TIMEOUT/NO STATUS surface. /// After the fix, disconnects count as attempts and exceed the cap. #[nativelink_test] -async fn worker_disconnect_loop_caps_at_max_job_retries_test() -> Result<(), Error> { +async fn worker_disconnect_loop_caps_at_max_worker_loss_retries_test() -> Result<(), Error> { let worker_id = WorkerId("worker_id".to_string()); let task_change_notify = Arc::new(Notify::new()); let (scheduler, _worker_scheduler) = SimpleScheduler::new_with_callback( &SimpleSpec { max_job_retries: 1, + max_worker_loss_retries: 1, ..Default::default() }, memory_awaited_action_db_factory( @@ -2485,7 +2489,7 @@ async fn worker_disconnect_loop_caps_at_max_job_retries_test() -> Result<(), Err OperationId::from(operation_id.as_str()) }; - // First disconnect: should requeue (attempts=1, not yet > max_job_retries=1). + // First loss: should requeue (worker_losses=1, not yet > max_worker_loss_retries=1). drop( scheduler .update_action( @@ -2519,9 +2523,9 @@ async fn worker_disconnect_loop_caps_at_max_job_retries_test() -> Result<(), Err ); } - // Second disconnect: now attempts=2 > max_job_retries=1, so the action - // must transition to Completed with an error mentioning the disconnect - // loop, not silently requeue. + // Second loss: now worker_losses=2 > max_worker_loss_retries=1, so the + // action must transition to Completed with an error saying the worker + // keeps being lost, not silently requeue. max_job_retries is untouched. drop( scheduler .update_action( @@ -2545,7 +2549,7 @@ async fn worker_disconnect_loop_caps_at_max_job_retries_test() -> Result<(), Err .expect("Completed action from disconnect cap must carry an error"); assert!( err.to_string() - .contains("Worker disconnected repeatedly while executing this action"), + .contains("was lost 2 times, more than max_worker_loss_retries (1)"), "Error message did not mention disconnect loop: {err}", ); assert_eq!( @@ -3377,7 +3381,7 @@ async fn failed_final_update_does_not_leak_worker_capacity() -> Result<(), Error async fn setup_worker_holding_a_finished_operation() -> Result< ( Arc, - mpsc::UnboundedReceiver, + mpsc::Receiver, OperationId, Box, ), @@ -3584,3 +3588,451 @@ async fn live_worker_that_never_acknowledges_a_kill_is_evicted() -> Result<(), E Ok(()) } + +/// Wraps a real `AwaitedActionDb` and makes the matching pass read the +/// queue one action per permit on `gate`, so a test can do things between +/// two actions of the same pass. Only queued-range reads are gated. +#[derive(MetricsComponent)] +struct GatedQueueDb { + inner: A, + gate: Arc, + /// When set, a queued listing also waits for one permit here before it + /// reads the queue at all, so a test can queue several actions and be + /// sure one pass lists them all. + listing_gate: Option>, +} + +impl AwaitedActionDb for GatedQueueDb { + type Subscriber = A::Subscriber; + + async fn get_awaited_action_by_id( + &self, + client_operation_id: &OperationId, + ) -> Result, Error> { + self.inner + .get_awaited_action_by_id(client_operation_id) + .await + } + + async fn get_all_awaited_actions( + &self, + ) -> Result> + Send, Error> { + self.inner.get_all_awaited_actions().await + } + + async fn get_by_operation_id( + &self, + operation_id: &OperationId, + ) -> Result, Error> { + self.inner.get_by_operation_id(operation_id).await + } + + async fn get_range_of_actions( + &self, + state: SortedAwaitedActionState, + start: Bound, + end: Bound, + desc: bool, + ) -> Result> + Send, Error> { + if let (Some(listing_gate), SortedAwaitedActionState::Queued) = (&self.listing_gate, state) + { + listing_gate.acquire().await.expect("gate closed").forget(); + } + let items = self + .inner + .get_range_of_actions(state, start, end, desc) + .await? + .collect::>() + .await; + let gate = matches!(state, SortedAwaitedActionState::Queued).then(|| self.gate.clone()); + Ok(futures::stream::iter(items).then(move |item| { + let gate = gate.clone(); + async move { + if let Some(gate) = gate { + gate.acquire_owned().await.expect("gate closed").forget(); + } + item + } + })) + } + + async fn update_awaited_action(&self, new_awaited_action: AwaitedAction) -> Result<(), Error> { + self.inner.update_awaited_action(new_awaited_action).await + } + + async fn add_action( + &self, + client_operation_id: OperationId, + action_info: Arc, + no_event_action_timeout: Duration, + ) -> Result { + self.inner + .add_action(client_operation_id, action_info, no_event_action_timeout) + .await + } +} + +fn start_execute_digest(update: UpdateForWorker) -> (String, DigestInfo) { + match update.update { + Some(update_for_worker::Update::StartAction(start_execute)) => { + let digest = start_execute + .execute_request + .and_then(|request| request.action_digest) + .and_then(|digest| DigestInfo::try_from(digest).ok()) + .expect("StartExecute carries an action digest"); + (start_execute.operation_id, digest) + } + v => panic!("Expected StartAction, got : {v:?}"), + } +} + +/// A worker that frees up while the pass is past the head of the queue +/// gets the oldest waiting action, not the one the pass happens to be +/// reading. Three actions, one worker with one slot: the first takes the +/// slot; the pass reads the second, which waits; the worker finishes; the +/// pass reads the third, and the second must be the one dispatched. +#[nativelink_test] +async fn freed_capacity_mid_pass_goes_to_the_oldest_waiting_action_test() -> Result<(), Error> { + let worker_id = WorkerId("worker_id".to_string()); + let task_change_notify = Arc::new(Notify::new()); + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let (scheduler, _worker_scheduler) = SimpleScheduler::new_with_callback( + &SimpleSpec::default(), + GatedQueueDb { + inner: memory_awaited_action_db_factory( + 0, + &task_change_notify.clone(), + MockInstantWrapped::default, + ), + gate: gate.clone(), + listing_gate: None, + }, + || async move {}, + task_change_notify, + MockInstantWrapped::default, + None, + ); + + let (tx, mut rx_from_worker) = mpsc::channel(64); + let worker = Worker::new( + worker_id.clone(), + PlatformProperties::default(), + tx, + NOW_TIME, + 1, + ); + scheduler.add_worker(worker).await?; + tokio::task::yield_now().await; + verify_initial_connection_message(worker_id.clone(), &mut rx_from_worker).await; + + let digests = [ + DigestInfo::new([1u8; 32], 512), + DigestInfo::new([2u8; 32], 512), + DigestInfo::new([3u8; 32], 512), + ]; + // The listeners stay alive so the actions keep their clients. + let _listener_a = + setup_action(&scheduler, digests[0], HashMap::new(), make_system_time(1)).await?; + let _listener_b = + setup_action(&scheduler, digests[1], HashMap::new(), make_system_time(2)).await?; + let _listener_c = + setup_action(&scheduler, digests[2], HashMap::new(), make_system_time(3)).await?; + + // The pass reads the first action: it takes the only slot. + gate.add_permits(1); + tokio::task::yield_now().await; + let (first_operation_id, first_digest) = + start_execute_digest(rx_from_worker.recv().await.unwrap()); + assert_eq!(first_digest, digests[0]); + + // The pass reads the second: the worker is full, so it waits. + gate.add_permits(1); + tokio::task::yield_now().await; + assert!( + poll!(Box::pin(rx_from_worker.recv())).is_pending(), + "nothing should be dispatched while the worker is full" + ); + + // The worker finishes the first action while the pass is still running. + let action_result = ActionResult { + output_files: Vec::new(), + output_folders: Vec::new(), + output_file_symlinks: Vec::new(), + output_directory_symlinks: Vec::new(), + exit_code: 0, + stdout_digest: DigestInfo::new([6u8; 32], 19), + stderr_digest: DigestInfo::new([7u8; 32], 20), + execution_metadata: ExecutionMetadata { + worker: worker_id.to_string(), + queued_timestamp: make_system_time(5), + worker_start_timestamp: make_system_time(6), + worker_completed_timestamp: make_system_time(7), + input_fetch_start_timestamp: make_system_time(8), + input_fetch_completed_timestamp: make_system_time(9), + execution_start_timestamp: make_system_time(10), + execution_completed_timestamp: make_system_time(11), + output_upload_start_timestamp: make_system_time(12), + output_upload_completed_timestamp: make_system_time(13), + }, + server_logs: HashMap::default(), + error: None, + message: String::new(), + }; + scheduler + .update_action( + &worker_id, + &OperationId::from(first_operation_id), + UpdateOperationType::UpdateWithActionStage(ActionStage::Completed(action_result)), + ) + .await?; + + // The pass reads the third action and finds room. The room goes to + // the second action, which has waited longer. + gate.add_permits(1); + tokio::task::yield_now().await; + let (_, dispatched) = start_execute_digest(rx_from_worker.recv().await.unwrap()); + assert_eq!( + dispatched, digests[1], + "the oldest waiting action takes the freed slot" + ); + + Ok(()) +} + +/// A completed result for `worker_id`, for tests that finish an action to +/// free its slot. +fn completed_result(worker_id: &WorkerId) -> ActionResult { + ActionResult { + output_files: Vec::new(), + output_folders: Vec::new(), + output_file_symlinks: Vec::new(), + output_directory_symlinks: Vec::new(), + exit_code: 0, + stdout_digest: DigestInfo::new([6u8; 32], 19), + stderr_digest: DigestInfo::new([7u8; 32], 20), + execution_metadata: ExecutionMetadata { + worker: worker_id.to_string(), + queued_timestamp: make_system_time(5), + worker_start_timestamp: make_system_time(6), + worker_completed_timestamp: make_system_time(7), + input_fetch_start_timestamp: make_system_time(8), + input_fetch_completed_timestamp: make_system_time(9), + execution_start_timestamp: make_system_time(10), + execution_completed_timestamp: make_system_time(11), + output_upload_start_timestamp: make_system_time(12), + output_upload_completed_timestamp: make_system_time(13), + }, + server_logs: HashMap::default(), + error: None, + message: String::new(), + } +} + +/// Parked actions are offered room in the order the queue listed them, not +/// shape by shape. Two actions fill a two-slot worker; three more of two +/// shapes park in listing order A1, B1, A2; both slots free up mid-pass; +/// A1 and B1 must take them, not A1 and A2. +#[nativelink_test] +async fn room_opening_mid_pass_keeps_the_listing_order_across_shapes() -> Result<(), Error> { + let worker_id = WorkerId("worker_id".to_string()); + let task_change_notify = Arc::new(Notify::new()); + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let (scheduler, _worker_scheduler) = SimpleScheduler::new_with_callback( + &SimpleSpec { + supported_platform_properties: Some(HashMap::from([( + "cpu_count".to_string(), + PropertyType::Minimum, + )])), + ..SimpleSpec::default() + }, + GatedQueueDb { + inner: memory_awaited_action_db_factory( + 0, + &task_change_notify.clone(), + MockInstantWrapped::default, + ), + gate: gate.clone(), + listing_gate: None, + }, + || async move {}, + task_change_notify, + MockInstantWrapped::default, + None, + ); + + let (tx, mut rx_from_worker) = mpsc::channel(64); + let worker = Worker::new( + worker_id.clone(), + PlatformProperties::new(HashMap::from([( + "cpu_count".to_string(), + PlatformPropertyValue::Minimum(8), + )])), + tx, + NOW_TIME, + 2, + ); + scheduler.add_worker(worker).await?; + tokio::task::yield_now().await; + verify_initial_connection_message(worker_id.clone(), &mut rx_from_worker).await; + + let cpu = |n: &str| HashMap::from([("cpu_count".to_string(), n.to_string())]); + let digests: Vec = (1..=6u8).map(|i| DigestInfo::new([i; 32], 512)).collect(); + // Listing order: two fillers, then A1 (cpu 1), B1 (cpu 2), A2 (cpu 1), + // then one more to read after the room opens. + let shapes = ["1", "1", "1", "2", "1", "1"]; + // The listeners stay alive so the actions keep their clients. + let _listeners = + futures::future::try_join_all(digests.iter().enumerate().map(|(i, digest)| { + setup_action( + &scheduler, + *digest, + cpu(shapes[i]), + make_system_time(u64::try_from(i).unwrap() + 1), + ) + })) + .await?; + + // The two fillers take the two slots. + gate.add_permits(2); + tokio::task::yield_now().await; + let (first_id, first) = start_execute_digest(rx_from_worker.recv().await.unwrap()); + let (second_id, second) = start_execute_digest(rx_from_worker.recv().await.unwrap()); + assert_eq!((first, second), (digests[0], digests[1])); + + // A1, B1 and A2 are read while the worker is full and park. + gate.add_permits(3); + tokio::task::yield_now().await; + assert!(poll!(Box::pin(rx_from_worker.recv())).is_pending()); + + // Both fillers finish while the pass is still running. + for id in [first_id, second_id] { + scheduler + .update_action( + &worker_id, + &OperationId::from(id), + UpdateOperationType::UpdateWithActionStage(ActionStage::Completed( + completed_result(&worker_id), + )), + ) + .await?; + } + + // The pass reads the sixth action and sees the room. It goes to A1 then + // B1, the two oldest parked, whatever their shapes. + gate.add_permits(1); + tokio::task::yield_now().await; + let (_, third) = start_execute_digest(rx_from_worker.recv().await.unwrap()); + let (_, fourth) = start_execute_digest(rx_from_worker.recv().await.unwrap()); + assert_eq!( + (third, fourth), + (digests[2], digests[3]), + "room goes to the parked actions in listing order, across shapes" + ); + assert!( + poll!(Box::pin(rx_from_worker.recv())).is_pending(), + "the worker is full again" + ); + Ok(()) +} + +/// A keepalive reporting more free memory is room opening, the same as an +/// action finishing: with the live memory veto on, the older action parked +/// for lack of reported memory takes it, not the newer one the pass reads +/// next. +#[nativelink_test] +async fn a_higher_memory_report_mid_pass_goes_to_the_oldest_waiting_action() -> Result<(), Error> { + let worker_id = WorkerId("worker_id".to_string()); + let task_change_notify = Arc::new(Notify::new()); + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let listing_gate = Arc::new(tokio::sync::Semaphore::new(0)); + let (scheduler, _worker_scheduler) = SimpleScheduler::new_with_callback( + &SimpleSpec { + supported_platform_properties: Some(HashMap::from([( + "memory_kb".to_string(), + PropertyType::Minimum, + )])), + live_memory_veto: Some("memory_kb".to_string()), + ..SimpleSpec::default() + }, + GatedQueueDb { + inner: memory_awaited_action_db_factory( + 0, + &task_change_notify.clone(), + MockInstantWrapped::default, + ), + gate: gate.clone(), + listing_gate: Some(listing_gate.clone()), + }, + || async move {}, + task_change_notify, + MockInstantWrapped::default, + None, + ); + + let (tx, mut rx_from_worker) = mpsc::channel(64); + let worker = Worker::new( + worker_id.clone(), + PlatformProperties::new(HashMap::from([( + "memory_kb".to_string(), + PlatformPropertyValue::Minimum(100_000), + )])), + tx, + NOW_TIME, + 1, + ); + scheduler.add_worker(worker).await?; + tokio::task::yield_now().await; + verify_initial_connection_message(worker_id.clone(), &mut rx_from_worker).await; + // The worker says it has almost nothing free. + scheduler + .worker_keep_alive_received( + &worker_id, + NOW_TIME + 1, + Some(WorkerLoad { + free_memory_kb: 1_000, + }), + ) + .await?; + + let memory = HashMap::from([("memory_kb".to_string(), "2000".to_string())]); + let digests = [ + DigestInfo::new([1u8; 32], 512), + DigestInfo::new([2u8; 32], 512), + ]; + let _listener_a = + setup_action(&scheduler, digests[0], memory.clone(), make_system_time(1)).await?; + let _listener_b = setup_action(&scheduler, digests[1], memory, make_system_time(2)).await?; + + // Both are queued before the pass lists the queue, so one pass reads + // both. It reads the older action first: vetoed for memory, it parks. + listing_gate.add_permits(1); + gate.add_permits(1); + for _ in 0..4 { + tokio::task::yield_now().await; + } + assert!(poll!(Box::pin(rx_from_worker.recv())).is_pending()); + + // The resident process exits and the next keepalive says so. + scheduler + .worker_keep_alive_received( + &worker_id, + NOW_TIME + 2, + Some(WorkerLoad { + free_memory_kb: 50_000, + }), + ) + .await?; + + // The pass reads the newer action and finds room; the room is the + // older action's. + gate.add_permits(1); + for _ in 0..4 { + tokio::task::yield_now().await; + } + let (_, dispatched) = start_execute_digest(rx_from_worker.recv().await.unwrap()); + assert_eq!( + dispatched, digests[0], + "the oldest waiting action takes the room a memory report opened" + ); + Ok(()) +} diff --git a/nativelink-scheduler/tests/store_awaited_action_db_test.rs b/nativelink-scheduler/tests/store_awaited_action_db_test.rs index efb1227a3ce..430a3fdb410 100644 --- a/nativelink-scheduler/tests/store_awaited_action_db_test.rs +++ b/nativelink-scheduler/tests/store_awaited_action_db_test.rs @@ -1,5 +1,6 @@ #![allow(clippy::todo)] +use core::ops::Bound; use core::sync::atomic::{AtomicUsize, Ordering}; use core::time::Duration; use std::collections::HashMap; @@ -7,12 +8,13 @@ use std::sync::Arc; use std::time::SystemTime; use bytes::Bytes; -use futures::{Stream, stream}; +use futures::{Stream, StreamExt, stream}; use mock_instant::thread_local::{MockClock, SystemTime as MockSystemTime}; use nativelink_error::Error; use nativelink_macro::nativelink_test; use nativelink_scheduler::awaited_action_db::{ - AwaitedAction, AwaitedActionDb, PersistedSortKey, SortedAwaitedAction, + AwaitedAction, AwaitedActionDb, AwaitedActionSubscriber, PersistedSortKey, SortedAwaitedAction, + SortedAwaitedActionState, }; use nativelink_scheduler::store_awaited_action_db::{ StoreAwaitedActionDb, inner_update_awaited_action, @@ -157,7 +159,7 @@ async fn test_inner_update_awaited_action() -> Result<(), Error> { assert_eq!( update.0, Bytes::from( - "{\"version\":0,\"action_info\":{\"command_digest\":\"0101010101010101010101010101010101010101010101010101010101010101-10\",\"input_root_digest\":\"0202020202020202020202020202020202020202020202020202020202020202-10\",\"timeout\":{\"secs\":1,\"nanos\":0},\"platform_properties\":{},\"priority\":0,\"load_timestamp\":{\"secs_since_epoch\":0,\"nanos_since_epoch\":0},\"insert_timestamp\":{\"secs_since_epoch\":0,\"nanos_since_epoch\":0},\"unique_qualifier\":{\"Uncacheable\":{\"instance_name\":\"foo\",\"digest_function\":\"Sha256\",\"digest\":\"0303030303030303030303030303030303030303030303030303030303030303-10\"}}},\"operation_id\":{\"String\":\"DEMO_OPERATION_ID\"},\"sort_key\":9223372041149743103,\"last_worker_updated_timestamp\":{\"secs_since_epoch\":0,\"nanos_since_epoch\":0},\"last_client_keepalive_timestamp\":{\"secs_since_epoch\":0,\"nanos_since_epoch\":0},\"worker_id\":null,\"state\":{\"stage\":\"Queued\",\"last_transition_timestamp\":{\"secs_since_epoch\":0,\"nanos_since_epoch\":0},\"client_operation_id\":{\"String\":\"DEMO_OPERATION_ID\"},\"action_digest\":\"0303030303030303030303030303030303030303030303030303030303030303-10\"},\"maybe_origin_metadata\":null,\"attempts\":0}" + "{\"version\":0,\"action_info\":{\"command_digest\":\"0101010101010101010101010101010101010101010101010101010101010101-10\",\"input_root_digest\":\"0202020202020202020202020202020202020202020202020202020202020202-10\",\"timeout\":{\"secs\":1,\"nanos\":0},\"platform_properties\":{},\"priority\":0,\"load_timestamp\":{\"secs_since_epoch\":0,\"nanos_since_epoch\":0},\"insert_timestamp\":{\"secs_since_epoch\":0,\"nanos_since_epoch\":0},\"unique_qualifier\":{\"Uncacheable\":{\"instance_name\":\"foo\",\"digest_function\":\"Sha256\",\"digest\":\"0303030303030303030303030303030303030303030303030303030303030303-10\"}}},\"operation_id\":{\"String\":\"DEMO_OPERATION_ID\"},\"sort_key\":9223372041149743103,\"last_worker_updated_timestamp\":{\"secs_since_epoch\":0,\"nanos_since_epoch\":0},\"last_client_keepalive_timestamp\":{\"secs_since_epoch\":0,\"nanos_since_epoch\":0},\"worker_id\":null,\"state\":{\"stage\":\"Queued\",\"last_transition_timestamp\":{\"secs_since_epoch\":0,\"nanos_since_epoch\":0},\"client_operation_id\":{\"String\":\"DEMO_OPERATION_ID\"},\"action_digest\":\"0303030303030303030303030303030303030303030303030303030303030303-10\"},\"maybe_origin_metadata\":null,\"attempts\":0,\"worker_losses\":0,\"escalations\":0}" ), "{update:#?}" ); @@ -697,3 +699,153 @@ async fn index_keys_from_both_versions_serve_older_entries_first() -> Result<(), } Ok(()) } + +/// Fake `SchedulerStore` whose listing returns the given actions and which +/// counts every `get_and_decode` by key prefix, so a test can see what a +/// listed subscriber reads when borrowed. +struct ListingStore { + encoded_actions: Vec, + reads_by_prefix: Mutex>, +} + +impl ListingStore { + fn new(actions: &[AwaitedAction]) -> Self { + Self { + encoded_actions: actions + .iter() + .map(|action| Bytes::from(serde_json::to_vec(action).expect("serialize"))) + .collect(), + reads_by_prefix: Mutex::new(HashMap::new()), + } + } +} + +impl SchedulerStore for ListingStore { + type SubscriptionManager = PendingSubscriptionManager; + + fn subscription_manager( + &self, + ) -> impl Future, Error>> { + std::future::ready(Ok(Arc::new(PendingSubscriptionManager))) + } + + fn update_data( + &self, + _data: T, + _expiry: Option, + ) -> impl Future, Error>> + where + T: SchedulerStoreDataProvider + + SchedulerStoreKeyProvider + + SchedulerCurrentVersionProvider + + Send, + { + std::future::ready(Ok(Some(1))) + } + + fn search_by_index_prefix( + &self, + _index: K, + ) -> impl Future< + Output = Result< + impl Stream::DecodeOutput, Error>> + Send, + Error, + >, + > + where + K: SchedulerIndexProvider + SchedulerStoreDecodeTo + Send, + ::DecodeOutput: Send, + { + let items: Vec<_> = self + .encoded_actions + .iter() + .map(|encoded| K::decode(1, encoded.clone())) + .collect(); + std::future::ready(Ok(stream::iter(items))) + } + + async fn count_by_index_prefix(&self, _index: K) -> Result + where + K: SchedulerIndexProvider + Send, + { + Ok(self.encoded_actions.len() as u64) + } + + async fn get_and_decode( + &self, + key: K, + ) -> Result::DecodeOutput>, Error> + where + K: SchedulerStoreKeyProvider + SchedulerStoreDecodeTo + Send, + { + let key = key.get_key().as_str().to_string(); + let prefix = key.split('_').next().unwrap_or("").to_string(); + *self.reads_by_prefix.lock().await.entry(prefix).or_default() += 1; + Ok(None) + } +} + +/// A listing already loads every record it returns, so borrowing a listed +/// subscriber must not read the record again. Only the client keepalive, +/// kept under its own key, is read, and once per subscriber. +#[nativelink_test] +async fn listed_subscriber_borrows_without_reading_the_record_again() -> Result<(), Error> { + fn new_op_id() -> OperationId { + OperationId::from("new-operation") + } + let actions: Vec = (0..3) + .map(|i| { + AwaitedAction::new( + OperationId::from(format!("op-{i}")), + make_cacheable_action_info(), + MockSystemTime::now().into(), + ) + }) + .collect(); + let store = Arc::new(ListingStore::new(&actions)); + let now_fn: fn() -> MockInstantWrapped = MockInstantWrapped::default; + let op_id_fn: fn() -> OperationId = new_op_id; + let db = StoreAwaitedActionDb::new( + store.clone(), + Arc::new(Notify::new()), + now_fn, + op_id_fn, + 60, + 60, + false, + ) + .await?; + + let subscribers: Vec<_> = db + .get_range_of_actions( + SortedAwaitedActionState::Queued, + Bound::Unbounded, + Bound::Unbounded, + true, + ) + .await? + .collect::>() + .await; + assert_eq!(subscribers.len(), 3); + for (subscriber, action) in subscribers.iter().zip(&actions) { + let subscriber = subscriber.as_ref().map_err(Clone::clone)?; + // Borrowed twice, as the matcher does. + assert_eq!( + subscriber.borrow().await?.operation_id(), + action.operation_id() + ); + assert_eq!( + subscriber.borrow().await?.operation_id(), + action.operation_id() + ); + } + + let reads = store.reads_by_prefix.lock().await; + assert_eq!(reads.get("aa"), None, "the record is not read again"); + assert_eq!( + reads.get("ck"), + Some(&3), + "the keepalive is read once per subscriber" + ); + Ok(()) +} diff --git a/nativelink-scheduler/tests/unsatisfiable_action_test.rs b/nativelink-scheduler/tests/unsatisfiable_action_test.rs index 8c64892946b..381daec1e73 100644 --- a/nativelink-scheduler/tests/unsatisfiable_action_test.rs +++ b/nativelink-scheduler/tests/unsatisfiable_action_test.rs @@ -134,8 +134,8 @@ async fn add_worker( name: &str, properties: PlatformProperties, max_inflight_tasks: u64, -) -> Result, Error> { - let (tx, rx) = mpsc::unbounded_channel(); +) -> Result, Error> { + let (tx, rx) = mpsc::channel(64); let worker = Worker::new( WorkerId(name.to_string()), properties, diff --git a/nativelink-scheduler/tests/utils/scheduler_utils.rs b/nativelink-scheduler/tests/utils/scheduler_utils.rs index bd14f638316..7036f3c59d9 100644 --- a/nativelink-scheduler/tests/utils/scheduler_utils.rs +++ b/nativelink-scheduler/tests/utils/scheduler_utils.rs @@ -78,7 +78,7 @@ pub(crate) async fn verify_overlapping_invocations( ("unscoped", None), ] { let worker_id = WorkerId(format!("fresh-worker-{invocation}")); - let (tx, mut rx) = mpsc::unbounded_channel(); + let (tx, mut rx) = mpsc::channel(64); // One in-flight action per worker, matching the single-use deployment. scheduler .add_worker(Worker::new( diff --git a/nativelink-service/src/health_server.rs b/nativelink-service/src/health_server.rs index 50c42e1558e..b72d012c63c 100644 --- a/nativelink-service/src/health_server.rs +++ b/nativelink-service/src/health_server.rs @@ -25,6 +25,7 @@ use http_body_util::Full; use hyper::header::{CONTENT_TYPE, HeaderValue}; use hyper::{Request, Response, StatusCode}; use nativelink_config::cas_server::HealthConfig; +use nativelink_error::{Error, make_input_err}; use nativelink_util::health_utils::{ HealthRegistry, HealthStatus, HealthStatusDescription, HealthStatusReporter, }; @@ -42,6 +43,36 @@ const DEFAULT_HEALTH_CHECK_TIMEOUT_SECONDS: u64 = 5; pub struct HealthServer { health_registry: HealthRegistry, timeout: Duration, + /// Readiness: a component still initializing makes the answer 503. + strict: bool, +} + +/// Where the plain status check answers when `path` is unset. +pub const DEFAULT_STATUS_PATH: &str = "/status"; +/// Where the readiness check answers when `readiness_path` is unset. +pub const DEFAULT_READINESS_PATH: &str = "/ready"; + +/// The two paths the health service serves, defaults filled in, as +/// `(status, readiness)`. They must differ: the router panics on a +/// duplicate route, so a configuration that names the same path for both +/// is refused here with a message instead. +pub fn health_paths(health_cfg: &HealthConfig) -> Result<(String, String), Error> { + let status = if health_cfg.path.is_empty() { + DEFAULT_STATUS_PATH + } else { + &health_cfg.path + }; + let readiness = if health_cfg.readiness_path.is_empty() { + DEFAULT_READINESS_PATH + } else { + &health_cfg.readiness_path + }; + if status == readiness { + return Err(make_input_err!( + "services.health.path and readiness_path are both {status}; the readiness check needs a path of its own" + )); + } + Ok((status.to_string(), readiness.to_string())) } impl HealthServer { @@ -54,8 +85,17 @@ impl HealthServer { Self { health_registry, timeout, + strict: false, } } + + /// The readiness form: unavailable while anything is initializing, not + /// only when something failed. + pub const fn readiness(health_registry: HealthRegistry, health_cfg: &HealthConfig) -> Self { + let mut server = Self::new(health_registry, health_cfg); + server.strict = true; + server + } } impl Service> for HealthServer { @@ -70,6 +110,7 @@ impl Service> for HealthServer { fn call(&mut self, _req: Request) -> Self::Future { let health_registry = self.health_registry.clone(); let local_timeout = self.timeout; + let strict = self.strict; Box::pin(error_span!("health_server_call").in_scope(|| async move { let health_status_descriptions: Vec = health_registry .health_status_report(&local_timeout) @@ -82,6 +123,11 @@ impl Service> for HealthServer { health_status_descriptions.iter().any(|description| { matches!(description.status, HealthStatus::Failed { .. }) | matches!(description.status, HealthStatus::Timeout { .. }) + | (strict + && matches!( + description.status, + HealthStatus::Initializing { .. } + )) }); let status_code = if contains_failed_report { StatusCode::SERVICE_UNAVAILABLE diff --git a/nativelink-service/src/worker_api_server.rs b/nativelink-service/src/worker_api_server.rs index 7414829ca69..86239a37256 100644 --- a/nativelink-service/src/worker_api_server.rs +++ b/nativelink-service/src/worker_api_server.rs @@ -28,7 +28,7 @@ use nativelink_proto::com::github::trace_machina::nativelink::remote_execution:: WorkerApi, WorkerApiServer as Server, }; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{ - execute_result, ExecuteComplete, ExecuteResult, GoingAwayRequest, KeepAliveRequest, UpdateForScheduler, UpdateForWorker + execute_declined, execute_result, ExecuteAccepted, ExecuteComplete, ExecuteDeclined, ExecuteResult, GoingAwayRequest, KeepAliveRequest, UpdateForScheduler, UpdateForWorker }; use nativelink_scheduler::worker::Worker; use nativelink_scheduler::worker_scheduler::WorkerScheduler; @@ -185,7 +185,9 @@ impl WorkerApiServer { )); }; - let (tx, rx) = mpsc::unbounded_channel(); + let (tx, rx) = mpsc::channel(nativelink_scheduler::worker::channel_capacity( + connect_worker_request.max_inflight_tasks, + )); // First convert our proto platform properties into one our scheduler understands. let platform_properties = { @@ -324,7 +326,9 @@ impl WorkerConnection { instance.inner_keep_alive(keep_alive_request).await } Update::GoingAwayRequest(going_away_request) => { - had_going_away = true; + // A drain keeps the worker until its stream closes, + // so that close still has to remove it. + had_going_away = !going_away_request.drain; instance.inner_going_away(going_away_request).await } Update::ExecuteResult(execute_result) => { @@ -333,6 +337,12 @@ impl WorkerConnection { Update::ExecuteComplete(execute_complete) => { instance.execution_complete(execute_complete).await } + Update::ExecuteAccepted(execute_accepted) => { + instance.dispatch_accepted(execute_accepted).await + } + Update::ExecuteDeclined(execute_declined) => { + instance.dispatch_declined(execute_declined).await + } }; if let Err(err) = result { tracing::warn!(worker_id=?instance.worker_id, ?err, "Error processing worker message"); @@ -362,7 +372,14 @@ impl WorkerConnection { Ok(()) } - async fn inner_going_away(&self, _going_away_request: GoingAwayRequest) -> Result<(), Error> { + async fn inner_going_away(&self, going_away_request: GoingAwayRequest) -> Result<(), Error> { + if going_away_request.drain { + self.scheduler + .set_drain_worker(&self.worker_id, true) + .await + .err_tip(|| "While draining worker in WorkerApiServer::inner_going_away")?; + return Ok(()); + } self.scheduler .remove_worker(&self.worker_id) .await @@ -371,10 +388,12 @@ impl WorkerConnection { } /// Any message from the worker proves it is alive, not only keepalives. - /// Only a keepalive carries a load report; the last one stands. + /// Only a keepalive carries a load report and lifts a pause; this + /// refreshes the timestamp alone, so a decline cannot undo the pause + /// it just took. async fn touch_liveness(&self) -> Result<(), Error> { self.scheduler - .worker_keep_alive_received(&self.worker_id, (self.now_fn)()?.as_secs(), None) + .worker_liveness_refreshed(&self.worker_id, (self.now_fn)()?.as_secs()) .await .err_tip(|| "Could not refresh worker liveness") } @@ -423,6 +442,36 @@ impl WorkerConnection { Ok(()) } + /// The answer is recorded before the liveness refresh: the refresh runs + /// the unacknowledged sweep, which would otherwise requeue an + /// acknowledgement that arrives right at the timeout. + async fn dispatch_accepted(&self, execute_accepted: ExecuteAccepted) -> Result<(), Error> { + let operation_id = OperationId::from(execute_accepted.operation_id); + self.scheduler + .worker_dispatch_accepted(&self.worker_id, &operation_id) + .await + .err_tip(|| format!("Failed to record acceptance of operation {operation_id}"))?; + self.touch_liveness().await + } + + async fn dispatch_declined(&self, execute_declined: ExecuteDeclined) -> Result<(), Error> { + let operation_id = OperationId::from(execute_declined.operation_id); + let reason = execute_declined::Reason::try_from(execute_declined.reason) + .unwrap_or(execute_declined::Reason::Unspecified); + let needs_kb = (reason == execute_declined::Reason::Load && execute_declined.needed_kb > 0) + .then_some(execute_declined.needed_kb); + let mut why = reason.as_str_name().to_ascii_lowercase(); + if !execute_declined.detail.is_empty() { + why.push_str(": "); + why.push_str(&execute_declined.detail); + } + self.scheduler + .worker_dispatch_declined(&self.worker_id, &operation_id, why, needs_kb) + .await + .err_tip(|| format!("Failed to record decline of operation {operation_id}"))?; + self.touch_liveness().await + } + async fn execution_complete(&self, execute_complete: ExecuteComplete) -> Result<(), Error> { self.touch_liveness().await?; let operation_id = OperationId::from(execute_complete.operation_id); diff --git a/nativelink-service/tests/health_server_test.rs b/nativelink-service/tests/health_server_test.rs index 85ae1fe19da..9b63515d747 100644 --- a/nativelink-service/tests/health_server_test.rs +++ b/nativelink-service/tests/health_server_test.rs @@ -6,7 +6,7 @@ use axum::http::Request; use hyper::StatusCode; use nativelink_config::cas_server::HealthConfig; use nativelink_macro::nativelink_test; -use nativelink_service::health_server::HealthServer; +use nativelink_service::health_server::{HealthServer, health_paths}; use nativelink_util::health_utils::{ HealthRegistry, HealthRegistryBuilder, HealthStatus, HealthStatusIndicator, }; @@ -155,3 +155,83 @@ async fn health_test_with_sleep() -> Result<(), Box> { )); Ok(()) } + +struct InitializingIndicator {} + +#[async_trait] +impl HealthStatusIndicator for InitializingIndicator { + fn get_name(&self) -> &'static str { + "initializing_indicator" + } + + async fn check_health(&self, _namespace: Cow<'static, str>) -> HealthStatus { + HealthStatus::Initializing { + struct_name: "InitializingIndicator", + message: "not yet".into(), + } + } + + fn struct_name(&self) -> &'static str { + "InitializingIndicator" + } +} + +async fn status_of(server: HealthServer) -> Result> { + let tonic_services = Routes::builder().routes(); + let mut svc = tonic_services + .into_axum_router() + .route_service("/probe", server); + let request = Request::builder() + .method("GET") + .uri("/probe") + .body(Body::empty())?; + let response: hyper::Response = + svc.as_service().ready().await?.call(request).await?; + Ok(response.status()) +} + +/// A component still initializing is healthy enough for `/status` and not +/// ready enough for `/ready`. +#[nativelink_test] +async fn readiness_is_unavailable_while_initializing() -> Result<(), Box> { + let mut health_registry_builder = HealthRegistryBuilder::new("foo"); + health_registry_builder.register_indicator(Arc::new(InitializingIndicator {})); + let health_registry = health_registry_builder.build(); + let config = HealthConfig::default(); + + assert_eq!( + status_of(HealthServer::new(health_registry.clone(), &config)).await?, + StatusCode::OK + ); + assert_eq!( + status_of(HealthServer::readiness(health_registry, &config)).await?, + StatusCode::SERVICE_UNAVAILABLE + ); + Ok(()) +} + +/// Defaults fill the two paths; a configuration that gives both the same +/// path is refused with a message rather than left to panic the router. +#[nativelink_test] +async fn health_paths_are_distinct_or_refused() -> Result<(), Box> { + assert_eq!( + health_paths(&HealthConfig::default())?, + ("/status".to_string(), "/ready".to_string()) + ); + let custom = HealthConfig { + path: "/healthz".to_string(), + readiness_path: "/readyz".to_string(), + ..HealthConfig::default() + }; + assert_eq!( + health_paths(&custom)?, + ("/healthz".to_string(), "/readyz".to_string()) + ); + let clash = HealthConfig { + path: "/ready".to_string(), + ..HealthConfig::default() + }; + let err = health_paths(&clash).expect_err("the same path for both must be refused"); + assert!(err.to_string().contains("both /ready"), "{err}"); + Ok(()) +} diff --git a/nativelink-service/tests/worker_api_server_test.rs b/nativelink-service/tests/worker_api_server_test.rs index c73028f648c..4b63d264380 100644 --- a/nativelink-service/tests/worker_api_server_test.rs +++ b/nativelink-service/tests/worker_api_server_test.rs @@ -13,7 +13,7 @@ // limitations under the License. use core::time::Duration; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::{Arc, Mutex}; use std::time::{SystemTime, UNIX_EPOCH}; @@ -31,14 +31,14 @@ use nativelink_proto::build::bazel::remote::execution::v2::{ }; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::update_for_scheduler::Update; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{ - execute_result, update_for_worker, ConnectWorkerRequest, ExecuteComplete, ExecuteResult, KeepAliveRequest, UpdateForScheduler + execute_declined, execute_result, update_for_worker, ConnectWorkerRequest, ExecuteAccepted, ExecuteComplete, ExecuteDeclined, ExecuteResult, GoingAwayRequest, KeepAliveRequest, UpdateForScheduler, WorkerLoad }; use nativelink_proto::google::rpc::Status as ProtoStatus; use nativelink_scheduler::api_worker_scheduler::ApiWorkerScheduler; use nativelink_scheduler::match_outcome::MatchOutcome; use nativelink_scheduler::platform_property_manager::PlatformPropertyManager; -use nativelink_scheduler::worker::ActionInfoWithProps; -use nativelink_scheduler::worker_scheduler::WorkerScheduler; +use nativelink_scheduler::worker::{ActionInfoWithProps, channel_capacity}; +use nativelink_scheduler::worker_scheduler::{WorkerScheduler, WorkerSummary}; use nativelink_service::worker_api_server::{ConnectWorkerStream, NowFn, WorkerApiServer}; use nativelink_util::action_messages::{ ActionInfo, ActionUniqueKey, ActionUniqueQualifier, OperationId, WorkerId, @@ -73,6 +73,8 @@ struct MockWorkerStateManager { tx_call: mpsc::UnboundedSender, rx_resp: Arc>>, tx_resp: mpsc::UnboundedSender, + /// Operations the state manager no longer has executing on any worker. + revoked: Mutex>, } impl MockWorkerStateManager { @@ -84,9 +86,15 @@ impl MockWorkerStateManager { tx_call, rx_resp: Arc::new(AsyncMutex::new(rx_resp)), tx_resp, + revoked: Mutex::new(HashSet::new()), } } + /// From now on the operation is not executing on any worker. + pub(crate) fn revoke(&self, operation_id: &OperationId) { + self.revoked.lock().unwrap().insert(operation_id.clone()); + } + pub(crate) async fn expect_update_operation( &self, result: Result<(), Error>, @@ -129,10 +137,10 @@ impl WorkerStateManager for MockWorkerStateManager { async fn is_executing_on_worker( &self, - _operation_id: &OperationId, + operation_id: &OperationId, _worker_id: &WorkerId, ) -> Result { - Ok(true) + Ok(!self.revoked.lock().unwrap().contains(operation_id)) } } @@ -161,6 +169,15 @@ async fn setup_api_server_with_task_limit( worker_timeout: u64, now_fn: NowFn, max_worker_tasks: u64, +) -> Result { + setup_api_server_with(worker_timeout, now_fn, max_worker_tasks, 0).await +} + +async fn setup_api_server_with( + worker_timeout: u64, + now_fn: NowFn, + max_worker_tasks: u64, + dispatch_ack_timeout_s: u64, ) -> Result { const SCHEDULER_NAME: &str = "DUMMY_SCHEDULE_NAME"; @@ -175,9 +192,11 @@ async fn setup_api_server_with_task_limit( platform_property_manager, WorkerAllocationStrategy::default(), None, + None, tasks_or_worker_change_notify, worker_timeout, 60, // unacknowledged_kill_timeout_s + dispatch_ack_timeout_s, worker_registry, None, false, // has_peers @@ -354,6 +373,7 @@ pub async fn server_does_not_timeout_if_execute_complete_test() origin_metadata: OriginMetadata::default(), scheduler_start_execute_event_id: None, }, + BASE_NOW_S, ) .await .unwrap(); @@ -393,8 +413,9 @@ pub async fn server_does_not_timeout_if_execute_complete_test() ); // And it is still evicted once it goes quiet for a full window. The - // eviction fails the running action through the state manager, so - // that call has to be serviced alongside it. + // eviction hands the running action back to the queue as a lost worker + // through the state manager, so that call has to be serviced alongside + // it. let timestamp = add_and_return_timestamp(BASE_WORKER_TIMEOUT_S); let (remove_result, (evicted_operation_id, evicted_worker_id, evicted_update)) = join!( test_context.scheduler.remove_timedout_workers(timestamp), @@ -404,8 +425,8 @@ pub async fn server_does_not_timeout_if_execute_complete_test() assert_eq!(evicted_operation_id, operation_id); assert_eq!(evicted_worker_id, test_context.worker_id); assert!( - matches!(evicted_update, UpdateOperationType::UpdateWithError(_)), - "expected the running action to be failed on eviction, got {evicted_update:?}" + matches!(evicted_update, UpdateOperationType::UpdateWithDisconnect), + "expected the running action to be requeued as a lost worker, got {evicted_update:?}" ); assert!( !test_context @@ -454,6 +475,7 @@ pub async fn stream_end_requeues_running_actions_as_disconnect_test() origin_metadata: OriginMetadata::default(), scheduler_start_execute_event_id: None, }, + BASE_NOW_S, ) .await .unwrap(); @@ -479,6 +501,49 @@ pub async fn stream_end_requeues_running_actions_as_disconnect_test() Ok(()) } +#[nativelink_test] +pub async fn going_away_with_drain_stops_dispatch_and_keeps_the_worker_until_the_stream_ends() +-> Result<(), Box> { + let test_context = setup_api_server(BASE_WORKER_TIMEOUT_S, Box::new(static_now_fn)).await?; + test_context + .worker_stream + .send(Update::GoingAwayRequest(GoingAwayRequest { drain: true })) + .await + .map_err(|e| make_err!(tonic::Code::Internal, "Error sending going away {e}"))?; + tokio::time::sleep(Duration::from_millis(10)).await; + + assert!( + test_context + .scheduler + .contains_worker_for_test(&test_context.worker_id) + .await, + "a draining worker stays in the pool while its stream is open" + ); + let outcome = test_context + .scheduler + .find_worker_for_action( + &PlatformProperties::new(HashMap::new()), + false, + make_system_time(BASE_NOW_S), + ) + .await; + assert!( + !matches!(outcome, MatchOutcome::Matched(_)), + "a draining worker must not be offered work: {outcome:?}" + ); + + drop(test_context.worker_stream); + tokio::time::sleep(Duration::from_millis(10)).await; + assert!( + !test_context + .scheduler + .contains_worker_for_test(&test_context.worker_id) + .await, + "the drained worker leaves the pool when its stream closes" + ); + Ok(()) +} + #[nativelink_test] pub async fn server_does_not_timeout_if_keep_alive_test() -> Result<(), Box> { @@ -642,6 +707,7 @@ pub async fn execution_response_success_test() -> Result<(), Box Result<(), Box Result<(), Box Result<(), Box> { + const ACK_TIMEOUT_S: u64 = 5; + let now_timestamp = Arc::new(Mutex::new(BASE_NOW_S)); + let now_timestamp_clone = now_timestamp.clone(); + let test_context = setup_api_server_with( + BASE_WORKER_TIMEOUT_S, + Box::new(move || Ok(Duration::from_secs(*now_timestamp_clone.lock().unwrap()))), + 0, + ACK_TIMEOUT_S, + ) + .await?; + + let action_digest = DigestInfo::new([8u8; 32], 123); + let action_info = Arc::new(ActionInfo { + command_digest: DigestInfo::new([0u8; 32], 0), + input_root_digest: DigestInfo::new([0u8; 32], 0), + timeout: Duration::MAX, + platform_properties: HashMap::new(), + priority: 0, + load_timestamp: make_system_time(0), + insert_timestamp: make_system_time(0), + unique_qualifier: ActionUniqueQualifier::Uncacheable(ActionUniqueKey { + execution_scope: None, + instance_name: "instance_name".to_string(), + digest_function: DigestHasherFunc::Sha256, + digest: action_digest, + }), + }); + let operation_id = OperationId::default(); + let platform_properties = test_context + .scheduler + .get_platform_property_manager() + .make_platform_properties(action_info.platform_properties.clone())?; + test_context + .scheduler + .worker_notify_run_action( + test_context.worker_id.clone(), + operation_id.clone(), + ActionInfoWithProps { + inner: action_info, + platform_properties, + origin_metadata: OriginMetadata::default(), + scheduler_start_execute_event_id: None, + }, + BASE_NOW_S, + ) + .await + .unwrap(); + + // The acknowledgement arrives past the timeout, with no message from + // the worker in between. + *now_timestamp.lock().unwrap() += ACK_TIMEOUT_S + 1; + test_context + .worker_stream + .send(Update::ExecuteAccepted(ExecuteAccepted { + operation_id: operation_id.to_string(), + })) + .await + .map_err(|e| make_err!(tonic::Code::Internal, "Error sending acceptance {e}"))?; + + // A requeue would go through the state manager; nothing may. + tokio::select! { + () = tokio::time::sleep(Duration::from_millis(50)) => {} + _ = test_context.state_manager.expect_update_operation(Ok(())) => { + panic!("the acknowledgement was requeued by the sweep it arrived with"); + } + } + assert!( + test_context + .scheduler + .running_action_info(&test_context.worker_id, &operation_id) + .await + .is_some(), + "the acknowledged dispatch stays on the worker" + ); + Ok(()) +} + +fn uncacheable_action(seed: u8) -> Arc { + Arc::new(ActionInfo { + command_digest: DigestInfo::new([0u8; 32], 0), + input_root_digest: DigestInfo::new([0u8; 32], 0), + timeout: Duration::MAX, + platform_properties: HashMap::new(), + priority: 0, + load_timestamp: make_system_time(0), + insert_timestamp: make_system_time(0), + unique_qualifier: ActionUniqueQualifier::Uncacheable(ActionUniqueKey { + execution_scope: None, + instance_name: "instance_name".to_string(), + digest_function: DigestHasherFunc::Sha256, + digest: DigestInfo::new([seed; 32], 123), + }), + }) +} + +/// Sends the action to the test worker as the matcher would, dated now. +async fn dispatch(test_context: &TestContext, seed: u8) -> Result { + let action_info = uncacheable_action(seed); + let platform_properties = test_context + .scheduler + .get_platform_property_manager() + .make_platform_properties(action_info.platform_properties.clone())?; + let operation_id = OperationId::default(); + test_context + .scheduler + .worker_notify_run_action( + test_context.worker_id.clone(), + operation_id.clone(), + ActionInfoWithProps { + inner: action_info, + platform_properties, + origin_metadata: OriginMetadata::default(), + scheduler_start_execute_event_id: None, + }, + BASE_NOW_S, + ) + .await?; + Ok(operation_id) +} + +/// The test worker's summary once the scheduler has taken in a keepalive +/// reporting `free_kb`: the worker's messages are handled on another task, +/// so the report is the sign that everything sent before it has been too. +async fn summary_after_report(test_context: &TestContext, free_kb: u64) -> WorkerSummary { + let worker_id = test_context.worker_id.to_string(); + for _ in 0..500 { + if let Some(summary) = test_context + .scheduler + .worker_snapshot() + .await + .into_iter() + .find(|summary| summary.id == worker_id && summary.free_memory_kb == Some(free_kb)) + { + return summary; + } + tokio::time::sleep(Duration::from_millis(2)).await; + } + panic!("the keepalive reporting {free_kb} KiB was never taken in"); +} + +/// A decline pauses the worker. The liveness refresh the decline carries +/// must not lift that pause, or the next matching pass hands the worker +/// the same action straight back; only a keepalive reporting enough room +/// resumes it. +#[nativelink_test] +pub async fn a_decline_keeps_the_worker_paused_until_a_keepalive_test() +-> Result<(), Box> { + let test_context = setup_api_server(BASE_WORKER_TIMEOUT_S, Box::new(static_now_fn)).await?; + let operation_id = dispatch(&test_context, 9).await?; + + // The decline sends the action back through the state manager. + let decline = test_context + .worker_stream + .send(Update::ExecuteDeclined(ExecuteDeclined { + operation_id: operation_id.to_string(), + reason: execute_declined::Reason::Load as i32, + detail: String::new(), + needed_kb: 4_000, + free_kb: 1_000, + })); + let (sent, (requeued_id, _, update)) = join!( + decline, + test_context.state_manager.expect_update_operation(Ok(())) + ); + sent.map_err(|e| make_err!(tonic::Code::Internal, "Error sending decline {e}"))?; + assert_eq!(requeued_id, operation_id); + assert!( + matches!(update, UpdateOperationType::UpdateWithDecline(_)), + "{update:?}" + ); + + // A keepalive without the room asked for leaves the pause in place. + test_context + .worker_stream + .send(Update::KeepAliveRequest(KeepAliveRequest { + load: Some(WorkerLoad { + free_memory_kb: 1_000, + }), + })) + .await + .map_err(|e| make_err!(tonic::Code::Internal, "Error sending keepalive {e}"))?; + let summary = summary_after_report(&test_context, 1_000).await; + assert!( + summary.is_paused, + "the decline's own refresh lifted the pause" + ); + + // One with the room lifts it. + test_context + .worker_stream + .send(Update::KeepAliveRequest(KeepAliveRequest { + load: Some(WorkerLoad { + free_memory_kb: 5_000, + }), + })) + .await + .map_err(|e| make_err!(tonic::Code::Internal, "Error sending keepalive {e}"))?; + let summary = summary_after_report(&test_context, 5_000).await; + assert!( + !summary.is_paused, + "a keepalive with the room resumes the worker" + ); + Ok(()) +} + +/// A kill that does not fit the worker's channel is a worker that is +/// behind, not one that is gone: it keeps its place and its actions, and +/// the next sweep sends the kill once the worker has read its backlog. +#[nativelink_test] +pub async fn a_full_channel_defers_the_kill_instead_of_evicting_test() +-> Result<(), Box> { + let mut test_context = + setup_api_server_with_task_limit(BASE_WORKER_TIMEOUT_S, Box::new(static_now_fn), 1).await?; + // Dispatches the worker never reads fill its channel to the brim. + let capacity = channel_capacity(1); + let mut operation_ids = Vec::new(); + for seed in 0..capacity { + operation_ids.push(dispatch(&test_context, u8::try_from(seed).unwrap()).await?); + } + let revoked = operation_ids[0].clone(); + test_context.state_manager.revoke(&revoked); + + // The kill has nowhere to go; the worker is neither evicted nor + // stripped of the operation. + test_context.scheduler.kill_revoked_operations().await?; + assert!( + test_context + .scheduler + .running_action_info(&test_context.worker_id, &revoked) + .await + .is_some(), + "a full channel evicted the worker" + ); + + // The worker reads one message; the next sweep gets the kill through. + let first = test_context + .connection_worker_stream + .next() + .await + .unwrap()?; + assert!(matches!( + first.update, + Some(update_for_worker::Update::StartAction(_)) + )); + test_context.scheduler.kill_revoked_operations().await?; + let mut deferred_kill = None; + for _ in 0..capacity { + let msg = test_context + .connection_worker_stream + .next() + .await + .unwrap()?; + if let Some(update_for_worker::Update::KillOperationRequest(kill)) = msg.update { + deferred_kill = Some(kill.operation_id); + break; + } + } + assert_eq!( + deferred_kill, + Some(revoked.to_string()), + "the deferred kill was never sent" + ); + Ok(()) +} diff --git a/nativelink-store/BUILD.bazel b/nativelink-store/BUILD.bazel index 91763e371a3..b82996e8816 100644 --- a/nativelink-store/BUILD.bazel +++ b/nativelink-store/BUILD.bazel @@ -57,6 +57,7 @@ rust_library( "src/redis_utils/ft_aggregate.rs", "src/redis_utils/ft_create.rs", "src/redis_utils/ft_cursor_read.rs", + "src/redis_utils/ft_info.rs", "src/redis_utils/ft_search_count.rs", "src/redis_utils/mod.rs", "src/ref_store.rs", diff --git a/nativelink-store/src/fast_slow_store.rs b/nativelink-store/src/fast_slow_store.rs index a6e9d28b153..299d432552f 100644 --- a/nativelink-store/src/fast_slow_store.rs +++ b/nativelink-store/src/fast_slow_store.rs @@ -467,6 +467,18 @@ impl StoreDriver for FastSlowStore { Ok(()) } + /// Both tiers drop their copy: the fast one is the cache this store + /// keeps, the slow one may be a cache of its own. + async fn remove(self: Pin<&Self>, key: StoreKey<'_>) -> Result { + let fast = self + .fast_store + .as_store_driver_pin() + .remove(key.borrow()) + .await?; + let slow = self.slow_store.as_store_driver_pin().remove(key).await?; + Ok(fast || slow) + } + async fn has_with_results( self: Pin<&Self>, key: &[StoreKey<'_>], diff --git a/nativelink-store/src/filesystem_store.rs b/nativelink-store/src/filesystem_store.rs index 0a0d00d187e..51459ca2acd 100644 --- a/nativelink-store/src/filesystem_store.rs +++ b/nativelink-store/src/filesystem_store.rs @@ -2055,6 +2055,10 @@ impl StoreDriver for FilesystemStore { Ok(()) } + async fn remove(self: Pin<&Self>, key: StoreKey<'_>) -> Result { + Ok(self.evicting_map.remove(&key.into_owned()).await) + } + async fn has_with_results( self: Pin<&Self>, keys: &[StoreKey<'_>], diff --git a/nativelink-store/src/grpc_store.rs b/nativelink-store/src/grpc_store.rs index 4ad8f9b661e..39ebff9923f 100644 --- a/nativelink-store/src/grpc_store.rs +++ b/nativelink-store/src/grpc_store.rs @@ -278,6 +278,13 @@ impl GrpcStore { spec: &GrpcSpec, jitter_fn: Arc Duration + Send + Sync>, ) -> Result, Error> { + error_if!( + matches!( + spec.store_type, + nativelink_config::stores::StoreType::EventSink + ) && spec.experimental_read_batching.is_some(), + "An event sink cannot batch reads" + ); error_if!( spec.endpoints.is_empty(), "Expected at least 1 endpoint in GrpcStore" @@ -1565,6 +1572,15 @@ impl StoreDriver for GrpcStore { keys: &[StoreKey<'_>], results: &mut [Option], ) -> Result<(), Error> { + if matches!( + self.store_type, + nativelink_config::stores::StoreType::EventSink + ) { + return Err(make_err!( + Code::Unimplemented, + "Event sinks do not support existence checks" + )); + } if matches!(self.store_type, nativelink_config::stores::StoreType::Ac) { keys.iter() .zip(results.iter_mut()) @@ -1629,7 +1645,7 @@ impl StoreDriver for GrpcStore { self: Pin<&Self>, key: StoreKey<'_>, reader: DropCloserReadHalf, - _size_info: UploadSizeInfo, + size_info: UploadSizeInfo, ) -> Result { struct LocalState { resource_name: String, @@ -1641,7 +1657,27 @@ impl StoreDriver for GrpcStore { } let is_digest_key = matches!(key, StoreKey::Digest(_)); - let digest = key.into_digest(); + let mut digest = key.into_digest(); + if matches!( + self.store_type, + nativelink_config::stores::StoreType::EventSink + ) { + let UploadSizeInfo::ExactSize(size) = size_info else { + return Err(make_err!( + Code::InvalidArgument, + "Event sinks require an exact upload size" + )); + }; + if is_digest_key { + return Err(make_err!( + Code::InvalidArgument, + "Event sinks require string keys" + )); + } + // String-key digests normally encode the key's length. Event bytes + // have a different length and must pass the stream size validator. + digest = DigestInfo::try_new(&digest.packed_hash().to_string(), size)?; + } if matches!(self.store_type, nativelink_config::stores::StoreType::Ac) { return self.update_action_result_from_bytes(digest, reader).await; } @@ -1745,13 +1781,27 @@ impl StoreDriver for GrpcStore { )) })); - self.write( - WriteRequestStreamWrapper::from(stream) - .await - .err_tip(|| "in GrpcStore::update()")?, - ) - .await - .err_tip(|| "in GrpcStore::update()")?; + let response = self + .write( + WriteRequestStreamWrapper::from(stream) + .await + .err_tip(|| "in GrpcStore::update()")?, + ) + .await + .err_tip(|| "in GrpcStore::update()")?; + + if matches!( + self.store_type, + nativelink_config::stores::StoreType::EventSink + ) && u64::try_from(response.get_ref().committed_size).ok() != Some(digest.size_bytes()) + { + return Err(make_err!( + Code::DataLoss, + "Event sink acknowledged {} bytes; expected {}", + response.get_ref().committed_size, + digest.size_bytes() + )); + } Ok(digest.size_bytes()) } @@ -1770,6 +1820,16 @@ impl StoreDriver for GrpcStore { read_limit: i64, } + if matches!( + self.store_type, + nativelink_config::stores::StoreType::EventSink + ) { + return Err(make_err!( + Code::Unimplemented, + "Event sinks do not support reads" + )); + } + let is_digest_key = matches!(key, StoreKey::Digest(_)); let digest = key.into_digest(); if matches!(self.store_type, nativelink_config::stores::StoreType::Ac) { diff --git a/nativelink-store/src/memory_store.rs b/nativelink-store/src/memory_store.rs index 70ffd27afa9..c107ad5f48a 100644 --- a/nativelink-store/src/memory_store.rs +++ b/nativelink-store/src/memory_store.rs @@ -103,6 +103,10 @@ impl StoreDriver for MemoryStore { Ok(()) } + async fn remove(self: Pin<&Self>, key: StoreKey<'_>) -> Result { + Ok(self.remove_entry(key).await) + } + async fn has_with_results( self: Pin<&Self>, keys: &[StoreKey<'_>], diff --git a/nativelink-store/src/redis_store.rs b/nativelink-store/src/redis_store.rs index b6d035cbbac..ef150ba6ccc 100644 --- a/nativelink-store/src/redis_store.rs +++ b/nativelink-store/src/redis_store.rs @@ -69,7 +69,7 @@ use uuid::Uuid; use crate::cas_utils::is_zero_digest; use crate::redis_utils::{ FtAggregateCursor, FtAggregateOptions, FtCreateOptions, SearchSchema, ft_aggregate, ft_create, - ft_search_count, + ft_info_indexing, ft_search_count, }; /// The default size of the read chunk when reading data from Redis. @@ -551,6 +551,34 @@ where }) } + /// `STRLEN` of a key, retried on a transient error like the reads are. + async fn strlen( + &self, + client: &mut ClientWithPermit, + encoded_key: &str, + ) -> Result { + let mut attempt: u32 = 0; + loop { + attempt += 1; + match redis::cmd("STRLEN") + .arg(encoded_key) + .query_async::(&mut client.connection_manager) + .await + { + Ok(v) => return Ok(v), + Err(err) + if attempt < MAX_REDIS_RETRY_ATTEMPTS && is_retryable_redis_error(&err) => + { + client.reconnect(&self.connection_manager).await?; + sleep(Duration::from_secs_f32(DEFAULT_RETRY_DELAY)).await; + } + Err(err) => { + return Err(Error::from(err).append("In RedisStore::get_part::strlen")); + } + } + } + } + async fn get_client(&self) -> Result, Error> { let local_client_permits = self.client_permits.clone(); let remaining = local_client_permits.available_permits(); @@ -1308,6 +1336,7 @@ where ); let mut client = self.get_client().await?; + let mut bytes_read: u64 = 0; loop { // getrange is position-based and idempotent, so re-resolve the // master and retry on a transient failover without re-sending @@ -1347,12 +1376,30 @@ where .send(chunk) .await .err_tip(|| "Failed to write data in RedisStore::get_part")?; + } else if bytes_read > 0 { + // GETRANGE answers the same for a value that ends on a + // chunk boundary and for a key removed between two + // chunks (eviction under allkeys-lru). Handing the second + // on as EOF is a truncated blob that every store above + // caches as complete, so ask how long the value is now + // and refuse anything but the end we stopped at. + let value_len = self.strlen(&mut client, encoded_key).await?; + let read_end = u64::try_from(data_start) + .unwrap_or(0) + .saturating_add(bytes_read); + if value_len != read_end { + return Err(make_err!( + Code::NotFound, + "Data for {key:?} changed under the read in the Redis store: read {bytes_read} bytes from offset {data_start}, the key now holds {value_len}" + )); + } } break; // No more data to read. } // We received a full chunk's worth of data, so write it... + bytes_read += chunk.len() as u64; writer .send(chunk) .await @@ -1624,6 +1671,42 @@ return {{ 1, new_version }} " ); +/// How long a freshly created index is given to finish its background +/// scan before a read goes ahead regardless, and how often it is asked. +const INDEX_READY_TIMEOUT: Duration = Duration::from_secs(5); +const INDEX_READY_POLL: Duration = Duration::from_millis(50); + +/// Polls `FT.INFO` until the index reports `indexing 0`, the timeout +/// passes, or the call fails; none of those stops the caller, which reads +/// the index either way. +async fn wait_for_index_ready(connection_manager: C, index: &str) +where + C: ConnectionLike + Send + Clone, +{ + let started = Instant::now(); + loop { + match ft_info_indexing(connection_manager.clone(), index).await { + Ok(false) => return, + Ok(true) => {} + Err(err) => { + warn!( + ?err, + index, "Could not read FT.INFO after creating the index; reading it now" + ); + return; + } + } + if started.elapsed() >= INDEX_READY_TIMEOUT { + warn!( + index, + "Index still backfilling after {INDEX_READY_TIMEOUT:?}; reading it now" + ); + return; + } + sleep(INDEX_READY_POLL).await; + } +} + /// This is the output of the calculations below hardcoded into the executable. const FINGERPRINT_CREATE_INDEX_HEX: &str = "3e762c15"; @@ -2283,6 +2366,24 @@ where } }); + // FT.CREATE returns before the index holds the keys that + // already exist; RediSearch scans them in the background. + // An aggregate in that window reads a partial index, and + // for the scheduler's queue that is a queue with actions + // missing, so wait for the scan to finish first. The wait + // is bounded; past it the aggregate runs anyway and the next + // pass re-reads. + if create_result.is_ok() { + wait_for_index_ready( + connection_manager.clone(), + &format!( + "{}", + get_index_name!(K::KEY_PREFIX, K::INDEX_NAME, K::MAYBE_SORT_KEY) + ), + ) + .await; + } + let run_result = run_ft_aggregate(connection_manager).await.err_tip(|| { format!( "Error with second ft_aggregate in RedisStore::search_by_index_prefix({})", diff --git a/nativelink-store/src/redis_utils/ft_aggregate.rs b/nativelink-store/src/redis_utils/ft_aggregate.rs index ab762dff006..f8e0f60a140 100644 --- a/nativelink-store/src/redis_utils/ft_aggregate.rs +++ b/nativelink-store/src/redis_utils/ft_aggregate.rs @@ -161,7 +161,7 @@ where )) } -fn resp2_data_parse( +pub(crate) fn resp2_data_parse( output: &mut RedisCursorData, results_array: &[Value], ) -> Result<(), RedisError> { @@ -208,7 +208,7 @@ fn resp2_data_parse( Ok(()) } -fn resp3_data_parse( +pub(crate) fn resp3_data_parse( output: &mut RedisCursorData, results_map: &Vec<(Value, Value)>, ) -> Result<(), RedisError> { diff --git a/nativelink-store/src/redis_utils/ft_cursor_read.rs b/nativelink-store/src/redis_utils/ft_cursor_read.rs index eb1323cf8a9..622576d8915 100644 --- a/nativelink-store/src/redis_utils/ft_cursor_read.rs +++ b/nativelink-store/src/redis_utils/ft_cursor_read.rs @@ -15,6 +15,7 @@ use redis::aio::ConnectionLike; use redis::{ErrorKind, RedisError, Value}; +use super::ft_aggregate::{resp2_data_parse, resp3_data_parse}; use crate::redis_utils::aggregate_types::RedisCursorData; pub(crate) async fn ft_cursor_read( @@ -41,14 +42,21 @@ where ))); } let mut value = value.into_iter(); - let Value::Array(data_ary) = value.next().unwrap() else { - return Err(RedisError::from((ErrorKind::Parse, "Non map item"))); - }; - if data_ary.is_empty() { - return Err(RedisError::from(( - ErrorKind::Parse, - "Expected at least 1 element in data array", - ))); + // The connection speaks RESP3, so a page past the first arrives as a + // map, the same shape the aggregate itself returns; RESP2 gives the + // bare array. Both parse the way the first page does, or every listing + // longer than one page failed here on its second page. + let mut output = RedisCursorData::default(); + match value.next().unwrap() { + Value::Array(data_ary) => resp2_data_parse(&mut output, &data_ary)?, + Value::Map(data_map) => resp3_data_parse(&mut output, &data_map)?, + other => { + return Err(RedisError::from(( + ErrorKind::Parse, + "Expected array or map for cursor results", + format!("{other:?}"), + ))); + } } let Value::Int(new_cursor_id) = value.next().unwrap() else { return Err(RedisError::from(( @@ -56,11 +64,6 @@ where "Expected cursor id as second element", ))); }; - - Ok(RedisCursorData { - // this should generally be impossible, but -1 provides a decent "obviously bad" value just in case - total: i64::try_from(data_ary.len()).unwrap_or(-1), - cursor: new_cursor_id as u64, - data: data_ary.into(), - }) + output.cursor = new_cursor_id as u64; + Ok(output) } diff --git a/nativelink-store/src/redis_utils/ft_info.rs b/nativelink-store/src/redis_utils/ft_info.rs new file mode 100644 index 00000000000..b3ad0a461b5 --- /dev/null +++ b/nativelink-store/src/redis_utils/ft_info.rs @@ -0,0 +1,78 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use redis::aio::ConnectionLike; +use redis::{ErrorKind, RedisError, Value}; + +/// Whether `RediSearch` is still backfilling `index`: `FT.INFO`'s +/// `indexing` field, `1` while the background scan that follows `FT.CREATE` +/// is running and `0` once every existing key is in the index. An aggregate +/// issued before that reads a partial index, which for the scheduler's queue +/// is a queue with actions missing from it. +pub(crate) async fn ft_info_indexing( + mut connection_manager: C, + index: &str, +) -> Result +where + C: ConnectionLike + Send, +{ + let value = redis::cmd("FT.INFO") + .arg(index) + .query_async::(&mut connection_manager) + .await?; + parse_indexing(&value) +} + +/// `FT.INFO` answers a flat key-value array on RESP2 and a map on RESP3. +fn parse_indexing(value: &Value) -> Result { + let field = |key: &Value, value: &Value| -> Option { + let name = match key { + Value::BulkString(bytes) => String::from_utf8_lossy(bytes).into_owned(), + Value::SimpleString(text) => text.clone(), + _ => return None, + }; + if name != "indexing" { + return None; + } + match value { + Value::Int(n) => Some(*n != 0), + Value::BulkString(bytes) => Some(bytes.as_slice() != b"0"), + Value::SimpleString(text) => Some(text != "0"), + _ => None, + } + }; + match value { + Value::Array(items) => { + for pair in items.chunks(2) { + if let [key, value] = pair + && let Some(indexing) = field(key, value) + { + return Ok(indexing); + } + } + } + Value::Map(entries) => { + for (key, value) in entries { + if let Some(indexing) = field(key, value) { + return Ok(indexing); + } + } + } + _ => {} + } + Err(RedisError::from(( + ErrorKind::Parse, + "FT.INFO reply carries no indexing field", + ))) +} diff --git a/nativelink-store/src/redis_utils/mod.rs b/nativelink-store/src/redis_utils/mod.rs index acdd5d0d7f3..3ceb7dc958c 100644 --- a/nativelink-store/src/redis_utils/mod.rs +++ b/nativelink-store/src/redis_utils/mod.rs @@ -16,7 +16,9 @@ mod aggregate_types; mod ft_aggregate; mod ft_create; mod ft_cursor_read; +mod ft_info; mod ft_search_count; pub(crate) use ft_aggregate::{FtAggregateCursor, FtAggregateOptions, ft_aggregate}; pub(crate) use ft_create::{FtCreateOptions, SearchSchema, ft_create}; +pub(crate) use ft_info::ft_info_indexing; pub(crate) use ft_search_count::ft_search_count; diff --git a/nativelink-store/src/ref_store.rs b/nativelink-store/src/ref_store.rs index 26c15678b9d..a44037658de 100644 --- a/nativelink-store/src/ref_store.rs +++ b/nativelink-store/src/ref_store.rs @@ -112,6 +112,10 @@ impl StoreDriver for RefStore { } } + async fn remove(self: Pin<&Self>, key: StoreKey<'_>) -> Result { + self.get_store()?.as_store_driver_pin().remove(key).await + } + async fn has_with_results( self: Pin<&Self>, keys: &[StoreKey<'_>], diff --git a/nativelink-store/src/size_partitioning_store.rs b/nativelink-store/src/size_partitioning_store.rs index b9883e42304..b66ac0f92a2 100644 --- a/nativelink-store/src/size_partitioning_store.rs +++ b/nativelink-store/src/size_partitioning_store.rs @@ -57,6 +57,19 @@ impl StoreDriver for SizePartitioningStore { Ok(()) } + async fn remove(self: Pin<&Self>, key: StoreKey<'_>) -> Result { + let StoreKey::Digest(digest) = key.borrow() else { + return Err(make_input_err!( + "SizePartitioningStore only supports Digest keys, got {key:?}" + )); + }; + if digest.size_bytes() < self.partition_size { + self.lower_store.as_store_driver_pin().remove(key).await + } else { + self.upper_store.as_store_driver_pin().remove(key).await + } + } + async fn has_with_results( self: Pin<&Self>, keys: &[StoreKey<'_>], diff --git a/nativelink-store/src/verify_store.rs b/nativelink-store/src/verify_store.rs index f42587aeb09..231ff7f641c 100644 --- a/nativelink-store/src/verify_store.rs +++ b/nativelink-store/src/verify_store.rs @@ -17,7 +17,7 @@ use std::sync::Arc; use async_trait::async_trait; use nativelink_config::stores::VerifySpec; -use nativelink_error::{Error, ResultExt, make_input_err}; +use nativelink_error::{Code, Error, ResultExt, make_err, make_input_err}; use nativelink_metric::MetricsComponent; use nativelink_util::buf_channel::{ DropCloserReadHalf, DropCloserWriteHalf, make_buf_channel_pair, @@ -31,6 +31,7 @@ use nativelink_util::metrics_utils::CounterWithTime; use nativelink_util::store_trait::{ RemoveCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo, }; +use tracing::warn; #[derive(Debug, MetricsComponent)] pub struct VerifyStore { @@ -151,6 +152,46 @@ impl VerifyStore { } } +impl VerifyStore { + async fn inner_check_read( + &self, + mut rx: DropCloserReadHalf, + writer: &mut DropCloserWriteHalf, + expected: u64, + key: StoreKey<'_>, + ) -> Result<(), Error> { + let mut sum_size: u64 = 0; + loop { + let chunk = rx + .recv() + .await + .err_tip(|| "Failed to read chunk in check_read in verify store")?; + if chunk.is_empty() { + if sum_size != expected { + self.size_verification_failures.inc(); + return Err(make_err!( + Code::DataLoss, + "Read {sum_size} bytes of the {expected} expected for {key:?} in verify store" + )); + } + return writer.send_eof().err_tip(|| "In verify_store::check_read"); + } + sum_size += chunk.len() as u64; + if sum_size > expected { + self.size_verification_failures.inc(); + return Err(make_err!( + Code::DataLoss, + "Read {sum_size} bytes, more than the {expected} expected for {key:?} in verify store" + )); + } + writer + .send(chunk) + .await + .err_tip(|| "Failed to forward chunk in verify store")?; + } + } +} + #[async_trait] impl StoreDriver for VerifyStore { async fn post_init(self: Arc) -> Result<(), Error> { @@ -230,7 +271,70 @@ impl StoreDriver for VerifyStore { offset: u64, length: Option, ) -> Result<(), Error> { - self.inner_store.get_part(key, writer, offset, length).await + // A digest names its size, so a read that ends short of it is a + // truncated blob, whatever tier below lost the tail: handed on as + // EOF it would be cached as complete by every store above and fed + // to actions as an input. The hash is not recomputed on reads. + let expected = match &key { + StoreKey::Digest(digest) if self.verify_size => { + let to_end = digest.size_bytes().saturating_sub(offset); + Some(length.map_or(to_end, |length| length.min(to_end))) + } + _ => None, + }; + let Some(expected) = expected else { + return self.inner_store.get_part(key, writer, offset, length).await; + }; + let (mut tx, rx) = make_buf_channel_pair(); + // The writer half lives inside the read future, so an inner read that + // fails after some data closes the channel as it returns; held out + // here it kept the forwarding loop waiting for bytes that never came. + let read_key = key.borrow().into_owned(); + let read_fut = async move { + self.inner_store + .get_part(read_key, &mut tx, offset, length) + .await + }; + let check_fut = self.inner_check_read(rx, writer, expected, key.borrow()); + let (read_res, check_res) = tokio::join!(read_fut, check_fut); + // The checker's verdict first: on a read past the expected size it + // drops its receiver and the inner send then fails with Internal, + // which must not hide the DataLoss. Otherwise the inner error (a + // NotFound from below) comes before the checker's closed-channel + // Internal, as `update` orders the same pair. + let result = match (read_res, check_res) { + (_, Err(e)) if e.code == Code::DataLoss => Err(e), + (Err(e), _) | (Ok(()), Err(e)) => Err(e), + (Ok(()), Ok(())) => Ok(()), + }; + if let Err(err) = &result + && err.code == Code::DataLoss + { + // The copy that served this read is bad; drop it from every + // cache below so the next read repopulates from the durable + // tier, or fails with NotFound and the client re-uploads. + match self + .inner_store + .as_store_driver_pin() + .remove(key.borrow()) + .await + { + Ok(removed) => warn!( + ?key, + removed, "Dropped a copy that failed size verification" + ), + Err(remove_err) => warn!( + ?key, + ?remove_err, + "Could not drop a copy that failed size verification" + ), + } + } + result + } + + async fn remove(self: Pin<&Self>, key: StoreKey<'_>) -> Result { + self.inner_store.as_store_driver_pin().remove(key).await } fn inner_store(&self, _digest: Option) -> &'_ dyn StoreDriver { diff --git a/nativelink-store/tests/grpc_store_test.rs b/nativelink-store/tests/grpc_store_test.rs index 4abf2685f4b..c06b170c9ee 100644 --- a/nativelink-store/tests/grpc_store_test.rs +++ b/nativelink-store/tests/grpc_store_test.rs @@ -29,7 +29,7 @@ use nativelink_util::background_spawn; use nativelink_util::buf_channel::make_buf_channel_pair; use nativelink_util::common::DigestInfo; use nativelink_util::proto_stream_utils::WriteRequestStreamWrapper; -use nativelink_util::store_trait::{StoreLike, UploadSizeInfo}; +use nativelink_util::store_trait::{StoreKey, StoreLike, UploadSizeInfo}; use nativelink_util::telemetry::ClientHeaders; use opentelemetry::Context; use regex::Regex; @@ -99,6 +99,8 @@ struct FakeStreamServer { /// Record every `WriteRequest`, not just the first, so a test can assert /// on how the client chunked the stream. drain_all: bool, + committed_size_override: Option, + fail_after_commit: Arc>, } impl FakeStreamServer { @@ -107,6 +109,8 @@ impl FakeStreamServer { write_requests: Arc::new(Mutex::new(vec![])), read_requests: Arc::new(Mutex::new(vec![])), drain_all: false, + committed_size_override: None, + fail_after_commit: Arc::new(Mutex::new(false)), } } @@ -177,7 +181,14 @@ impl ByteStream for FakeStreamServer { committed_size += i64::try_from(req.data.len()).unwrap(); self.write_requests.lock().await.push(req); } - return Ok(Response::new(WriteResponse { committed_size })); + if core::mem::take(&mut *self.fail_after_commit.lock().await) { + return Err(Status::unavailable( + "simulated lost acknowledgement after commit", + )); + } + return Ok(Response::new(WriteResponse { + committed_size: self.committed_size_override.unwrap_or(committed_size), + })); } let write_request = match stream.next().await { None => { @@ -276,6 +287,112 @@ async fn write_update_works_with_legacy_resource_names() -> Result<(), Error> { write_update_works_core(true, upload_pattern).await } +#[nativelink_test] +async fn event_sink_uses_payload_size_and_rejects_cache_operations() -> Result<(), Error> { + let (server, port) = make_fake_bytestream_server_draining().await; + let mut spec = test_spec(format!("http://localhost:{port}"), false); + spec.store_type = StoreType::EventSink; + spec.instance_name = "bep".into(); + let store = GrpcStore::new(&spec)?; + let key = StoreKey::from("BepEvent:be:test:1"); + let payload = b"event payload with a length different from its storage key"; + store + .update_oneshot(key.borrow(), payload.as_slice().into()) + .await?; + let writes = server.write_requests.lock().await; + assert!(writes[0].resource_name.starts_with("bep/uploads/")); + assert!(writes[0].resource_name.ends_with(&format!( + "/{}/{}", + key.borrow().into_digest().packed_hash(), + payload.len() + ))); + assert_eq!( + writes + .iter() + .flat_map(|request| request.data.iter().copied()) + .collect::>(), + payload + ); + assert!(writes.last().unwrap().finish_write); + drop(writes); + assert_eq!( + store.has(key.borrow()).await.unwrap_err().code, + tonic::Code::Unimplemented + ); + let (tx, _rx) = make_buf_channel_pair(); + assert_eq!( + store + .get_part(key.borrow(), tx, 0, None) + .await + .unwrap_err() + .code, + tonic::Code::Unimplemented + ); + let digest = DigestInfo::try_new(VALID_HASH, payload.len())?; + assert_eq!( + store + .update_oneshot(digest, payload.as_slice().into()) + .await + .unwrap_err() + .code, + tonic::Code::InvalidArgument + ); + let (_tx, rx) = make_buf_channel_pair(); + assert_eq!( + store + .update(key, rx, UploadSizeInfo::MaxSize(100)) + .await + .unwrap_err() + .code, + tonic::Code::InvalidArgument + ); + Ok(()) +} + +#[nativelink_test] +async fn event_sink_rejects_incomplete_acknowledgement() -> Result<(), Error> { + let server = FakeStreamServer { + committed_size_override: Some(0), + ..FakeStreamServer::new_draining() + }; + let (_, port) = spawn_bytestream_server(server).await; + let mut spec = test_spec(format!("http://localhost:{port}"), false); + spec.store_type = StoreType::EventSink; + let store = GrpcStore::new(&spec)?; + let err = store + .update_oneshot(StoreKey::from("event-ack"), "event payload".into()) + .await + .unwrap_err(); + assert_eq!(err.code, tonic::Code::DataLoss); + Ok(()) +} + +#[nativelink_test] +async fn event_sink_retry_preserves_upload_identity_after_lost_ack() -> Result<(), Error> { + let server = FakeStreamServer::new_draining(); + *server.fail_after_commit.lock().await = true; + let (server, port) = spawn_bytestream_server(server).await; + let mut spec = test_spec(format!("http://localhost:{port}"), false); + spec.store_type = StoreType::EventSink; + spec.retry.max_retries = 1; + let store = GrpcStore::new(&spec)?; + let payload = b"event delivered before the acknowledgement was lost"; + store + .update_oneshot(StoreKey::from("event-retry"), payload.as_slice().into()) + .await?; + let requests = server.write_requests.lock().await; + let writes: Vec<_> = requests + .iter() + .filter(|request| !request.data.is_empty()) + .collect(); + assert_eq!(writes.len(), 2); + assert_eq!(writes[0].resource_name, writes[1].resource_name); + assert_eq!(writes[0].write_offset, writes[1].write_offset); + assert_eq!(writes[0].data, payload.as_slice()); + assert_eq!(writes[1].data, payload.as_slice()); + Ok(()) +} + async fn read_works_core( use_legacy_resource_names: bool, upload_pattern: &str, @@ -531,9 +648,9 @@ async fn update_splits_buffers_larger_than_the_grpc_message_limit() -> Result<() let (server, port) = make_fake_bytestream_server_draining().await; let mut spec = test_spec(format!("http://localhost:{port}"), false); // This test moves and reassembles almost 10 MiB under instrumented builds. - // Keep the RPC deadline comfortably above sanitizer overhead; timeout - // behavior is not what this test exercises. - spec.rpc_timeout_s = 5; + // The deadline is headroom for that, not part of the test: at 5 s the + // ASan job timed out on a loaded runner with nothing else wrong. + spec.rpc_timeout_s = 30; let store = GrpcStore::new(&spec)?; let digest = DigestInfo::try_new(VALID_HASH, BLOB_LEN).unwrap(); @@ -691,7 +808,8 @@ async fn batch_update_blobs_splits_an_oversized_batch_across_rpcs() -> Result<() let (server, port) = make_fake_cas_server().await; let mut spec = test_spec(format!("http://localhost:{port}"), false); spec.instance_name = "backend_instance".to_string(); - spec.rpc_timeout_s = 5; + // Over 10 MiB through the fake server; headroom for sanitizer builds. + spec.rpc_timeout_s = 30; let store = GrpcStore::new(&spec)?; let requests: Vec = (0..ENTRIES) diff --git a/nativelink-store/tests/redis_store_test.rs b/nativelink-store/tests/redis_store_test.rs index 58703ef7bcf..b668633a778 100644 --- a/nativelink-store/tests/redis_store_test.rs +++ b/nativelink-store/tests/redis_store_test.rs @@ -65,6 +65,20 @@ fn make_temp_key(final_name: &str) -> String { format!("temp-{TEMP_UUID}-{{{final_name}}}") } +/// What `FT.INFO` says while the index created a moment ago is still +/// scanning existing keys (`indexing 1`) and once it is done (`0`). +fn make_ft_info(index: &str, indexing: i64) -> MockCmd { + MockCmd::new( + redis::cmd("FT.INFO").arg(index), + Ok(Value::Array(vec![ + Value::BulkString(b"num_docs".to_vec()), + Value::Int(3), + Value::BulkString(b"indexing".to_vec()), + Value::Int(indexing), + ])), + ) +} + async fn make_mock_store( commands: Vec, ) -> RedisStore> { @@ -699,6 +713,79 @@ async fn zero_len_items_exist_check() -> Result<(), Error> { Ok(()) } +// The key was evicted between two GETRANGEs: the second comes back empty, +// which is also what the end of a value on a chunk boundary looks like. The +// store asks STRLEN, sees the key is gone, and fails the read instead of +// ending it early. +#[nativelink_test] +async fn get_part_key_removed_between_chunks_is_not_found() -> Result<(), Error> { + let digest = DigestInfo::try_new(VALID_HASH1, DEFAULT_READ_CHUNK_SIZE * 2)?; + let real_key = format!("{digest}"); + let chunk_end = i64::try_from(DEFAULT_READ_CHUNK_SIZE).unwrap() - 1; + + let commands = vec![ + MockCmd::new( + redis::cmd("GETRANGE") + .arg(real_key.clone()) + .arg(0) + .arg(chunk_end), + Ok(Value::BulkString(vec![b'a'; DEFAULT_READ_CHUNK_SIZE])), + ), + MockCmd::new( + redis::cmd("GETRANGE") + .arg(real_key.clone()) + .arg(chunk_end + 1) + .arg(chunk_end * 2 + 1), + Ok(Value::BulkString(vec![])), + ), + MockCmd::new(redis::cmd("STRLEN").arg(real_key), Ok(Value::Int(0))), + ]; + let store = make_mock_store(commands).await; + + let result = store.get_part_unchunked(digest, 0, None).await; + assert_eq!( + result.as_ref().unwrap_err().code, + Code::NotFound, + "{result:?}" + ); + Ok(()) +} + +// A value that ends exactly on a chunk boundary: the empty GETRANGE is the +// real end, STRLEN agrees, and the read completes with every byte. +#[nativelink_test] +async fn get_part_value_ending_on_a_chunk_boundary_reads_whole() -> Result<(), Error> { + let digest = DigestInfo::try_new(VALID_HASH1, DEFAULT_READ_CHUNK_SIZE)?; + let real_key = format!("{digest}"); + let chunk_end = i64::try_from(DEFAULT_READ_CHUNK_SIZE).unwrap() - 1; + + let commands = vec![ + MockCmd::new( + redis::cmd("GETRANGE") + .arg(real_key.clone()) + .arg(0) + .arg(chunk_end), + Ok(Value::BulkString(vec![b'a'; DEFAULT_READ_CHUNK_SIZE])), + ), + MockCmd::new( + redis::cmd("GETRANGE") + .arg(real_key.clone()) + .arg(chunk_end + 1) + .arg(chunk_end * 2 + 1), + Ok(Value::BulkString(vec![])), + ), + MockCmd::new( + redis::cmd("STRLEN").arg(real_key), + Ok(Value::Int(i64::try_from(DEFAULT_READ_CHUNK_SIZE).unwrap())), + ), + ]; + let store = make_mock_store(commands).await; + + let result = store.get_part_unchunked(digest, 0, None).await?; + assert_eq!(result.len(), DEFAULT_READ_CHUNK_SIZE); + Ok(()) +} + #[nativelink_test] async fn list_test() -> Result<(), Error> { async fn get_list( @@ -1336,6 +1423,7 @@ fn test_search_by_index() -> Result<(), Error> { .arg("TAG"), Ok(Value::Nil), ), + make_ft_info("test:_content_prefix__3e762c15", 0), make_ft_aggregate(), ]; let store = make_mock_store(commands).await; @@ -1616,6 +1704,10 @@ fn test_search_by_index_skips_docs_that_expired_mid_query() -> Result<(), Error> .arg("TAG"), Ok(Value::Nil), ), + // The index is still backfilling on the first ask and ready on the + // second; the read waits for the second. + make_ft_info("test:_content_prefix__3e762c15", 1), + make_ft_info("test:_content_prefix__3e762c15", 0), make_ft_aggregate(), ]; let store = make_mock_store(commands).await; @@ -1727,6 +1819,7 @@ fn test_search_by_index_swallows_already_exists_from_ft_create() -> Result<(), E let commands = vec![ make_ft_aggregate(), make_ft_create_already_exists(), + make_ft_info("test:_content_prefix_sort_key_3e762c15", 0), make_ft_aggregate(), ]; let store = make_mock_store(commands).await; @@ -1895,6 +1988,7 @@ fn test_search_by_index_with_sort_key() -> Result<(), Error> { .arg("TAG"), Ok(Value::Nil), ), + make_ft_info("test:_content_prefix__3e762c15", 0), make_ft_aggregate(), ]; let store = make_mock_store(commands).await; @@ -2005,6 +2099,7 @@ fn test_search_by_index_resp3() -> Result<(), Error> { .arg("SORTABLE"), Ok(Value::Nil), ), + make_ft_info("test:_content_prefix_sort_key_3e762c15", 0), make_ft_aggregate(), ]; let store = make_mock_store(commands).await; @@ -2117,6 +2212,113 @@ fn test_search_by_index_skips_int_from_cursor_read() -> Result<(), Error> { Ok(()) } +/// The store speaks RESP3, under which a cursor page arrives as a map in +/// the same shape as the aggregate's own first page. Before the cursor read +/// parsed that shape, every listing longer than one page failed on its +/// second page with "Non map item". +#[nativelink_test] +async fn search_by_index_prefix_reads_a_resp3_cursor_page() -> Result<(), Error> { + fn make_ft_aggregate() -> MockCmd { + MockCmd::new( + redis::cmd("FT.AGGREGATE") + .arg("test:_content_prefix_sort_key_3e762c15") + .arg("@content_prefix:{ Searchable }") + .arg("TIMEOUT") + .arg(10000_u64) + .arg("LOAD") + .arg(2) + .arg("data") + .arg("version") + .arg("WITHCURSOR") + .arg("COUNT") + .arg(1500) + .arg("MAXIDLE") + .arg(30000) + .arg("SORTBY") + .arg(2usize) + .arg("@sort_key") + .arg("ASC") + .arg("MAX") + .arg(1_000_000_u64), + Ok(Value::Array(vec![ + Value::Array(vec![ + Value::Int(2), + Value::Array(vec![ + Value::BulkString(b"data".to_vec()), + Value::BulkString(b"first".to_vec()), + Value::BulkString(b"version".to_vec()), + Value::BulkString(b"1".to_vec()), + ]), + ]), + Value::Int(42), + ])), + ) + } + + fn make_ft_cursor_read_resp3() -> MockCmd { + MockCmd::new( + redis::cmd("ft.cursor") + .arg("read") + .arg("test:_content_prefix_sort_key_3e762c15") + .cursor_arg(42), + Ok(Value::Array(vec![ + Value::Map(vec![ + ( + Value::SimpleString("attributes".to_string()), + Value::Array(vec![]), + ), + ( + Value::SimpleString("format".to_string()), + Value::SimpleString("STRING".to_string()), + ), + ( + Value::SimpleString("results".to_string()), + Value::Array(vec![Value::Map(vec![ + ( + Value::SimpleString("extra_attributes".to_string()), + Value::Map(vec![ + ( + Value::BulkString(b"data".to_vec()), + Value::BulkString(b"second".to_vec()), + ), + ( + Value::BulkString(b"version".to_vec()), + Value::BulkString(b"2".to_vec()), + ), + ]), + ), + ( + Value::SimpleString("values".to_string()), + Value::Array(vec![]), + ), + ])]), + ), + ]), + Value::Int(0), + ])), + ) + } + + let store = make_mock_store(vec![make_ft_aggregate(), make_ft_cursor_read_resp3()]).await; + let search_results: Vec = store + .search_by_index_prefix(SearchByContentPrefix { + prefix: "Searchable".to_string(), + }) + .await + .err_tip(|| "Failed to search by index")? + .try_collect() + .await?; + + assert_eq!( + search_results.len(), + 2, + "the RESP3 page is read like the first" + ); + assert_eq!(search_results[0].content, "first"); + assert_eq!(search_results[1].content, "second"); + Ok(()) +} + #[nativelink_test] async fn no_items_from_none_subscription_channel() -> Result<(), Error> { let (_tx, rx) = tokio::sync::mpsc::unbounded_channel(); diff --git a/nativelink-store/tests/verify_store_test.rs b/nativelink-store/tests/verify_store_test.rs index fc1aa92e4ea..fee11c1a9ff 100644 --- a/nativelink-store/tests/verify_store_test.rs +++ b/nativelink-store/tests/verify_store_test.rs @@ -17,7 +17,7 @@ use core::pin::Pin; use futures::future::pending; use futures::try_join; use nativelink_config::stores::{MemorySpec, StoreSpec, VerifySpec}; -use nativelink_error::{Error, ResultExt}; +use nativelink_error::{Code, Error, ResultExt}; use nativelink_macro::nativelink_test; use nativelink_store::memory_store::MemoryStore; use nativelink_store::verify_store::VerifyStore; @@ -383,3 +383,305 @@ async fn verify_size_and_hash_succeeds_on_small_data() -> Result<(), Error> { ); Ok(()) } + +// The tier below holds fewer bytes than the digest names: the read fails +// instead of ending early, so nothing above caches a truncated blob. +#[nativelink_test] +async fn verify_size_true_fails_a_short_read() -> Result<(), Error> { + let inner_store = MemoryStore::new(&MemorySpec::default()); + let store = VerifyStore::new( + &VerifySpec { + backend: StoreSpec::Memory(MemorySpec::default()), + verify_size: true, + verify_hash: false, + }, + Store::new(inner_store.clone()), + ); + // Ten bytes claimed, four stored: the inner store does not check. + let digest = DigestInfo::try_new(VALID_HASH1, 10).unwrap(); + inner_store.update_oneshot(digest, "1234".into()).await?; + + // A range read that stops before the end is judged against the range, + // and passes while the bytes are there. + let result = store.get_part_unchunked(digest, 1, Some(3)).await?; + assert_eq!(result, "234"); + + let result = store.get_part_unchunked(digest, 0, None).await; + let err = result.unwrap_err(); + assert_eq!(err.code, Code::DataLoss, "{err:?}"); + assert!( + err.to_string().contains("Read 4 bytes of the 10 expected"), + "{err:?}" + ); + // The failed read dropped the copy, so a range past its end now finds + // nothing at all. + let result = store.get_part_unchunked(digest, 1, Some(5)).await; + assert_eq!(result.unwrap_err().code, Code::NotFound); + Ok(()) +} + +#[nativelink_test] +async fn verify_size_true_passes_a_whole_read() -> Result<(), Error> { + let inner_store = MemoryStore::new(&MemorySpec::default()); + let store = VerifyStore::new( + &VerifySpec { + backend: StoreSpec::Memory(MemorySpec::default()), + verify_size: true, + verify_hash: false, + }, + Store::new(inner_store.clone()), + ); + let digest = DigestInfo::try_new(VALID_HASH1, 4).unwrap(); + inner_store.update_oneshot(digest, "1234".into()).await?; + + assert_eq!(store.get_part_unchunked(digest, 0, None).await?, "1234"); + assert_eq!(store.get_part_unchunked(digest, 2, None).await?, "34"); + assert_eq!(store.get_part_unchunked(digest, 1, Some(2)).await?, "23"); + // Past the end there is nothing to read and nothing expected. + assert_eq!(store.get_part_unchunked(digest, 4, None).await?, ""); + Ok(()) +} + +/// The tier below sends part of the blob and then fails, as the Redis store +/// does for a key evicted between two chunks: the read fails with that +/// error and does not wait for bytes that will never come. +#[nativelink_test] +async fn verify_size_true_fails_when_the_inner_read_errors_after_some_data() -> Result<(), Error> { + use core::pin::Pin; + use std::sync::Arc; + + use async_trait::async_trait; + use nativelink_error::{Code, make_err}; + use nativelink_metric::{ + MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent, + }; + use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; + use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator}; + use nativelink_util::store_trait::{RemoveItemCallback, StoreDriver, StoreKey}; + + #[derive(Debug)] + struct PartialThenErrorStore; + + impl MetricsComponent for PartialThenErrorStore { + fn publish( + &self, + _kind: MetricKind, + _field_metadata: MetricFieldData, + ) -> Result { + Ok(MetricPublishKnownKindData::Component) + } + } + + #[async_trait] + impl StoreDriver for PartialThenErrorStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + async fn has_with_results( + self: Pin<&Self>, + _keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + results.fill(Some(10)); + Ok(()) + } + async fn update( + self: Pin<&Self>, + _key: StoreKey<'_>, + _reader: DropCloserReadHalf, + _size_info: UploadSizeInfo, + ) -> Result { + Ok(0) + } + async fn get_part( + self: Pin<&Self>, + _key: StoreKey<'_>, + writer: &mut DropCloserWriteHalf, + _offset: u64, + _length: Option, + ) -> Result<(), Error> { + writer.send("1234".into()).await?; + Err(make_err!(Code::NotFound, "key removed under the read")) + } + fn inner_store(&self, _key: Option) -> &dyn StoreDriver { + self + } + fn as_any<'a>(&'a self) -> &'a (dyn core::any::Any + Sync + Send + 'static) { + self + } + fn as_any_arc(self: Arc) -> Arc { + self + } + fn register_remove_callback( + self: Arc, + _callback: Arc, + ) -> Result<(), Error> { + Ok(()) + } + } + default_health_status_indicator!(PartialThenErrorStore); + + let store = VerifyStore::new( + &VerifySpec { + backend: StoreSpec::Memory(MemorySpec::default()), + verify_size: true, + verify_hash: false, + }, + Store::new(Arc::new(PartialThenErrorStore)), + ); + let digest = DigestInfo::try_new(VALID_HASH1, 10).unwrap(); + let result = tokio::time::timeout( + core::time::Duration::from_secs(5), + store.get_part_unchunked(digest, 0, None), + ) + .await + .expect("a failed inner read must not hang the verify store"); + let err = result.unwrap_err(); + assert_eq!(err.code, Code::NotFound, "{err:?}"); + assert!( + err.to_string().contains("key removed under the read"), + "{err:?}" + ); + Ok(()) +} + +/// The tier below serves more bytes than the digest names: the checker's +/// `DataLoss` is what the caller sees, not the inner store's failed send +/// once the checker stops reading. +#[nativelink_test] +async fn verify_size_true_reports_data_loss_on_an_over_long_read() -> Result<(), Error> { + use core::pin::Pin; + use std::sync::Arc; + + use async_trait::async_trait; + use nativelink_error::Code; + use nativelink_metric::{ + MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent, + }; + use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; + use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator}; + use nativelink_util::store_trait::{RemoveItemCallback, StoreDriver, StoreKey}; + + /// Serves forty bytes for every key, in four chunks, whatever the digest + /// says. + #[derive(Debug)] + struct OverLongStore; + + impl MetricsComponent for OverLongStore { + fn publish( + &self, + _kind: MetricKind, + _field_metadata: MetricFieldData, + ) -> Result { + Ok(MetricPublishKnownKindData::Component) + } + } + + #[async_trait] + impl StoreDriver for OverLongStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + async fn has_with_results( + self: Pin<&Self>, + _keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + results.fill(Some(10)); + Ok(()) + } + async fn update( + self: Pin<&Self>, + _key: StoreKey<'_>, + _reader: DropCloserReadHalf, + _size_info: UploadSizeInfo, + ) -> Result { + Ok(0) + } + async fn get_part( + self: Pin<&Self>, + _key: StoreKey<'_>, + writer: &mut DropCloserWriteHalf, + _offset: u64, + _length: Option, + ) -> Result<(), Error> { + for _ in 0..4 { + writer.send("0123456789".into()).await?; + } + writer.send_eof() + } + fn inner_store(&self, _key: Option) -> &dyn StoreDriver { + self + } + fn as_any<'a>(&'a self) -> &'a (dyn core::any::Any + Sync + Send + 'static) { + self + } + fn as_any_arc(self: Arc) -> Arc { + self + } + fn register_remove_callback( + self: Arc, + _callback: Arc, + ) -> Result<(), Error> { + Ok(()) + } + } + default_health_status_indicator!(OverLongStore); + + let store = VerifyStore::new( + &VerifySpec { + backend: StoreSpec::Memory(MemorySpec::default()), + verify_size: true, + verify_hash: false, + }, + Store::new(Arc::new(OverLongStore)), + ); + let digest = DigestInfo::try_new(VALID_HASH1, 10).unwrap(); + let err = tokio::time::timeout( + core::time::Duration::from_secs(5), + store.get_part_unchunked(digest, 0, None), + ) + .await + .expect("an over-long read must not hang") + .unwrap_err(); + assert_eq!(err.code, Code::DataLoss, "{err:?}"); + assert!( + err.to_string().contains("more than the 10 expected"), + "{err:?}" + ); + Ok(()) +} + +/// A copy that fails size verification is dropped from the tier that holds +/// it, so the next read repopulates from below instead of failing the same +/// way until the cache happens to evict it. +#[nativelink_test] +async fn verify_size_true_drops_a_copy_that_fails_verification() -> Result<(), Error> { + use nativelink_error::Code; + + let inner_store = MemoryStore::new(&MemorySpec::default()); + let store = VerifyStore::new( + &VerifySpec { + backend: StoreSpec::Memory(MemorySpec::default()), + verify_size: true, + verify_hash: false, + }, + Store::new(inner_store.clone()), + ); + // Ten bytes claimed, four stored: a truncated copy. + let digest = DigestInfo::try_new(VALID_HASH1, 10).unwrap(); + inner_store.update_oneshot(digest, "1234".into()).await?; + assert!(inner_store.has(digest).await?.is_some()); + + let err = store.get_part_unchunked(digest, 0, None).await.unwrap_err(); + assert_eq!(err.code, Code::DataLoss, "{err:?}"); + assert_eq!( + inner_store.has(digest).await?, + None, + "the truncated copy is gone after the failed read" + ); + // And the next read reports NotFound, the signal for a re-upload. + let err = store.get_part_unchunked(digest, 0, None).await.unwrap_err(); + assert_eq!(err.code, Code::NotFound, "{err:?}"); + Ok(()) +} diff --git a/nativelink-util/BUILD.bazel b/nativelink-util/BUILD.bazel index 3e8c8fde076..43f61050daf 100644 --- a/nativelink-util/BUILD.bazel +++ b/nativelink-util/BUILD.bazel @@ -115,6 +115,7 @@ rust_test_suite( "tests/proto_stream_utils_test.rs", "tests/resource_info_test.rs", "tests/retry_test.rs", + "tests/shutdown_guard_test.rs", "tests/store_trait_test.rs", "tests/telemetry_test.rs", "tests/tls_utils_test.rs", diff --git a/nativelink-util/src/metrics.rs b/nativelink-util/src/metrics.rs index 571cc156025..86c1d90adee 100644 --- a/nativelink-util/src/metrics.rs +++ b/nativelink-util/src/metrics.rs @@ -32,6 +32,7 @@ pub const EXECUTION_RESULT: &str = "execution.result"; pub const EXECUTION_INSTANCE: &str = "execution.instance"; pub const EXECUTION_PRIORITY: &str = "execution.priority"; pub const EXECUTION_WORKER_ID: &str = "execution.worker_id"; +pub const EXECUTION_OUTCOME: &str = "execution.outcome"; pub const EXECUTION_ACTION_MNEMONIC: &str = "execution.action_mnemonic"; pub const EXECUTION_EXIT_CODE: &str = "execution.exit_code"; pub const EXECUTION_ACTION_DIGEST: &str = "execution.action_digest"; @@ -706,12 +707,17 @@ pub struct ExecutionMetrics { /// Wall time is already covered by `execution.stage.duration`. This is the /// other half: an action pinning eight cores for a minute and one sleeping /// for a minute look identical by wall time and nothing alike here. -pub fn record_execution_cpu_time(cpu_time_ms: u64, instance_name: &str, action_mnemonic: &str) { +pub fn record_execution_cpu_time( + cpu_time_ms: u64, + instance_name: &str, + action_mnemonic: &str, + outcome: &str, +) { #[expect(clippy::cast_precision_loss)] // Milliseconds; f64 is exact well past any real action. let seconds = cpu_time_ms as f64 / 1000.0; EXECUTION_METRICS.execution_cpu_time.record( seconds, - &execution_sample_attributes(instance_name, action_mnemonic), + &execution_sample_attributes(instance_name, action_mnemonic, outcome), ); } @@ -725,17 +731,22 @@ pub fn record_execution_peak_memory( peak_memory_kb: u64, instance_name: &str, action_mnemonic: &str, + outcome: &str, ) { EXECUTION_METRICS.execution_peak_memory.record( peak_memory_sample(peak_memory_kb), - &execution_sample_attributes(instance_name, action_mnemonic), + &execution_sample_attributes(instance_name, action_mnemonic, outcome), ); } /// Attributes for a per-action resource sample. The mnemonic is attached only /// when known: an empty value on every sample from an unattributed client /// would read as a real mnemonic named "". -fn execution_sample_attributes(instance_name: &str, action_mnemonic: &str) -> Vec { +fn execution_sample_attributes( + instance_name: &str, + action_mnemonic: &str, + outcome: &str, +) -> Vec { let mut attrs = vec![KeyValue::new(EXECUTION_INSTANCE, instance_name.to_string())]; if !action_mnemonic.is_empty() { attrs.push(KeyValue::new( @@ -743,6 +754,9 @@ fn execution_sample_attributes(instance_name: &str, action_mnemonic: &str) -> Ve action_mnemonic.to_string(), )); } + if !outcome.is_empty() { + attrs.push(KeyValue::new(EXECUTION_OUTCOME, outcome.to_string())); + } attrs } @@ -886,6 +900,11 @@ pub static WORKER_METRICS: LazyLock = LazyLock::new(|| { .with_unit("{worker}") .build(), + dispatch_requeues: meter + .u64_counter("scheduler.dispatch.requeues") + .with_description("Dispatches sent back to the queue untried, by reason: declined, channel_full, unacknowledged") + .build(), + worker_keepalives: meter .u64_counter("worker.keepalives") .with_description("Total worker keepalives received") @@ -928,6 +947,8 @@ pub struct WorkerMetrics { pub worker_connections: metrics::Counter, /// Workers that have left, cumulative, by reason. pub worker_disconnections: metrics::Counter, + /// Dispatches requeued untried, by reason. + pub dispatch_requeues: metrics::Counter, /// Keepalives received, cumulative. pub worker_keepalives: metrics::Counter, /// Seconds between consecutive messages from a worker. @@ -972,6 +993,14 @@ pub fn record_worker_state(state: &'static str, entered: bool) { ); } +/// Records a dispatch going back to the queue untried: the worker declined +/// it, its channel was full, or it never acknowledged it. +pub fn record_dispatch_requeue(reason: &'static str) { + WORKER_METRICS + .dispatch_requeues + .add(1, &[KeyValue::new("reason", reason)]); +} + /// Records a keepalive from a worker. pub fn record_worker_keepalive() { WORKER_METRICS.worker_keepalives.add(1, &[]); @@ -1091,6 +1120,14 @@ pub static SCHEDULER_METRICS: LazyLock = LazyLock::new(|| .with_unit("{action}") .build(), + parked_dispatched: meter + .u64_counter("scheduler.matching.parked_dispatched") + .with_description( + "Queued actions dispatched onto room that opened while the matching pass was past them, taken ahead of the newer actions the pass was reading", + ) + .with_unit("{action}") + .build(), + awaited_action_orphans: meter .u64_counter("scheduler.awaited_action.orphans") .with_description( @@ -1104,6 +1141,14 @@ pub static SCHEDULER_METRICS: LazyLock = LazyLock::new(|| .with_description("Abandoned-queue sweeps that ended in an error before finishing") .with_unit("{sweep}") .build(), + + hint_resolutions: meter + .u64_counter("scheduler.hints.resolutions") + .with_description( + "Actions the historical resource scheduler sized, by where the numbers came from: a hint keyed by target, action digest, command digest or mnemonic, the cold start, or nothing", + ) + .with_unit("{action}") + .build(), } }); @@ -1129,10 +1174,15 @@ pub struct SchedulerOtlpMetrics { pub queue_depth: metrics::Gauge, /// Queued actions retired for having no client. pub queue_retired: metrics::Counter, + /// Actions a pass parked for want of room and dispatched when room + /// opened before the pass ended, ahead of the newer actions it was reading. + pub parked_dispatched: metrics::Counter, /// Listed operations whose record was gone on read, by site. pub awaited_action_orphans: metrics::Counter, /// Abandoned-queue sweeps that failed before finishing. pub sweep_failures: metrics::Counter, + /// Actions sized by the historical resource scheduler, by source. + pub hint_resolutions: metrics::Counter, } /// Records a completed matching pass. @@ -1166,6 +1216,15 @@ pub fn record_queue_retired(count: u64) { } } +/// Records parked actions a pass sent to a worker onto room that opened +/// mid-pass: confirmed sends, not attempts. A dispatch that failed or lost +/// the assignment race to another scheduler is not counted. +pub fn record_parked_dispatched(count: u64) { + if count > 0 { + SCHEDULER_METRICS.parked_dispatched.add(count, &[]); + } +} + /// Records an operation the queue listed whose record was gone when read; /// `site` is where it was met (`matching`, `sweep`). pub fn record_awaited_action_orphan(site: &'static str) { @@ -1179,6 +1238,15 @@ pub fn record_sweep_failure() { SCHEDULER_METRICS.sweep_failures.add(1, &[]); } +/// Records where an action's reservation came from: `target`, +/// `action_digest`, `command_digest`, `mnemonic`, `cold_start` or `none`. +/// The cold-start share of a run is `cold_start` over the total. +pub fn record_hint_resolution(source: &'static str) { + SCHEDULER_METRICS + .hint_resolutions + .add(1, &[KeyValue::new("source", source)]); +} + /// Records a queued action failed for being unsatisfiable. `properties` /// names the ones no worker could satisfy, sorted and comma separated, so /// the label set is bounded by the declared platform properties. diff --git a/nativelink-util/src/operation_state_manager.rs b/nativelink-util/src/operation_state_manager.rs index fd26b7f0e72..4e6e2a30dee 100644 --- a/nativelink-util/src/operation_state_manager.rs +++ b/nativelink-util/src/operation_state_manager.rs @@ -53,6 +53,11 @@ pub trait ActionStateResult: Send + Sync + 'static { async fn changed(&mut self) -> Result<(Arc, Option), Error>; /// Provide result as action info. This behavior will not be supported by all implementations. async fn as_action_info(&self) -> Result<(Arc, Option), Error>; + /// When a client last checked in on this operation, where the + /// implementation tracks it. `None` where it does not. + async fn client_last_seen(&self) -> Result, Error> { + Ok(None) + } } /// The direction in which the results are ordered. @@ -134,6 +139,36 @@ pub enum UpdateOperationType { /// Notification that the execution stage has completed and it's just IO happening now. ExecutionComplete, + + /// The worker never ran the action: it declined the dispatch, the + /// dispatch never reached it, or it never acknowledged it. Requeue + /// without counting an attempt; nothing was tried. + UpdateWithDecline(Decline), + + /// The worker killed the action for a resource it had declared too + /// little of; requeue it with a larger reservation, within the retry cap. + UpdateWithEscalation(Escalation), +} + +/// Why a dispatched action is going back to the queue untried. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Decline { + pub reason: String, +} + +/// A reservation to raise before an action runs again. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Escalation { + /// The platform property carrying the reservation. + pub property: String, + /// The new value. + pub value: u64, + /// Why: the error the kill produced, kept for the client if the cap + /// ends the escalation. + pub reason: Error, + /// Set on the last step: the CPU property and the largest worker's + /// whole CPU, so the action runs alone there. + pub cpu: Option<(String, u64)>, } #[async_trait] diff --git a/nativelink-util/src/shutdown_guard.rs b/nativelink-util/src/shutdown_guard.rs index c0d79972af6..7b0ce25a227 100644 --- a/nativelink-util/src/shutdown_guard.rs +++ b/nativelink-util/src/shutdown_guard.rs @@ -94,7 +94,11 @@ impl Default for ShutdownGuard { fn default() -> Self { let priority = Priority::LeastImportant; let mut map = HashMap::new(); - map.insert(priority, 0); + // The original counts too. Left at zero, its first `wait_for` + // promotion subtracted one from the count its clones had built up, + // so a single clone (the worker draining on SIGTERM) looked already + // dropped and the process exited before the drain began. + map.insert(priority, 1); let (tx, rx) = watch::channel(map); Self { priority, tx, rx } } @@ -102,9 +106,11 @@ impl Default for ShutdownGuard { impl Clone for ShutdownGuard { fn clone(&self) -> Self { + // A clone always starts least important, so it is counted there, + // whatever the original has been promoted to. self.tx.send_modify(|map| { map.insert( - self.priority, + Priority::LeastImportant, map.get(&Priority::LeastImportant) .unwrap_or(&0) .saturating_add(1), diff --git a/nativelink-util/src/store_trait.rs b/nativelink-util/src/store_trait.rs index 0882da0e145..7f652c2f901 100644 --- a/nativelink-util/src/store_trait.rs +++ b/nativelink-util/src/store_trait.rs @@ -625,6 +625,14 @@ pub trait StoreDriver: // for ref stores async fn post_init(self: Arc) -> Result<(), Error>; + /// Drops this store's own copy of `key`, so the next read repopulates it + /// from below; a store that holds the durable copy keeps it and answers + /// `false`. Used when a read proves a cached copy is bad. Returns whether + /// anything was dropped. + async fn remove(self: Pin<&Self>, _key: StoreKey<'_>) -> Result { + Ok(false) + } + /// See: [`StoreLike::has`] for details. #[inline] async fn has(self: Pin<&Self>, key: StoreKey<'_>) -> Result, Error> { diff --git a/nativelink-util/tests/metrics_test.rs b/nativelink-util/tests/metrics_test.rs index c293b63d784..435e70dd235 100644 --- a/nativelink-util/tests/metrics_test.rs +++ b/nativelink-util/tests/metrics_test.rs @@ -350,8 +350,8 @@ fn test_new_metric_helpers_are_callable() { record_store_tier_read("slow", "hit"); record_store_tier_io("fast", "read", 4096); record_health_check("store", "ok"); - record_execution_peak_memory(512 * 1024, "main", "CppCompile"); - record_execution_cpu_time(90_000, "main", "CppCompile"); + record_execution_peak_memory(512 * 1024, "main", "CppCompile", "COMPLETED"); + record_execution_cpu_time(90_000, "main", "CppCompile", "COMPLETED"); record_connection_acquired("grpc", Some(12), false); record_connection_acquired("redis", Some(0), true); record_connection_acquired("grpc", None, false); diff --git a/nativelink-util/tests/shutdown_guard_test.rs b/nativelink-util/tests/shutdown_guard_test.rs new file mode 100644 index 00000000000..dd1a2a27fbf --- /dev/null +++ b/nativelink-util/tests/shutdown_guard_test.rs @@ -0,0 +1,65 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use core::time::Duration; + +use nativelink_macro::nativelink_test; +use nativelink_util::shutdown_guard::{Priority, ShutdownGuard}; +use tokio::sync::broadcast; + +/// The original's wait must hold while a clone is alive and return once +/// it is dropped. Before, the promotion inside `wait_for` took one off the +/// clone's count, so a single clone looked already gone. +#[nativelink_test] +async fn wait_for_holds_until_the_clone_drops() { + let mut original = ShutdownGuard::default(); + let clone = original.clone(); + + let held = + tokio::time::timeout(Duration::from_millis(50), original.wait_for(Priority::P0)).await; + assert!( + held.is_err(), + "wait_for returned while a clone was still alive" + ); + + drop(clone); + tokio::time::timeout(Duration::from_secs(1), original.wait_for(Priority::P0)) + .await + .expect("wait_for should return once the clone is dropped"); +} + +/// The production path: the SIGTERM handler broadcasts a clone, a listener +/// receives it (another clone), the buffered value is dropped, and the +/// handler waits. The listener's guard is what has to hold the process. +#[nativelink_test] +async fn a_guard_received_over_a_broadcast_holds_the_wait() { + let mut original = ShutdownGuard::default(); + let (tx, mut rx) = broadcast::channel::(4); + tx.send(original.clone()).expect("send"); + let listener_guard = rx.recv().await.expect("recv"); + drop(tx); + drop(rx); + + let held = + tokio::time::timeout(Duration::from_millis(50), original.wait_for(Priority::P0)).await; + assert!( + held.is_err(), + "wait_for returned while the listener still held its guard" + ); + + drop(listener_guard); + tokio::time::timeout(Duration::from_secs(1), original.wait_for(Priority::P0)) + .await + .expect("wait_for should return once the listener drops its guard"); +} diff --git a/nativelink-worker/BUILD.bazel b/nativelink-worker/BUILD.bazel index d9f9821d216..b32036527eb 100644 --- a/nativelink-worker/BUILD.bazel +++ b/nativelink-worker/BUILD.bazel @@ -20,6 +20,7 @@ rust_library( "src/persistent_worker/pool.rs", "src/persistent_worker/protocol.rs", "src/qos.rs", + "src/reaper.rs", "src/running_actions_manager.rs", "src/worker_api_client_wrapper.rs", "src/worker_utils.rs", @@ -73,6 +74,7 @@ rust_test_suite( "tests/directory_cache_test.rs", "tests/local_worker_test.rs", "tests/namespace_utils_test.rs", + "tests/reaper_test.rs", "tests/running_actions_manager_test.rs", "tests/worker_utils_test.rs", ], @@ -102,6 +104,7 @@ rust_test_suite( "@crates//:prost", "@crates//:prost-types", "@crates//:serial_test", + "@crates//:tempfile", "@crates//:tokio", "@crates//:tonic", "@crates//:tonic-prost", diff --git a/nativelink-worker/src/buck2_file_capture.rs b/nativelink-worker/src/buck2_file_capture.rs index 2d19738a832..9f69eead161 100644 --- a/nativelink-worker/src/buck2_file_capture.rs +++ b/nativelink-worker/src/buck2_file_capture.rs @@ -31,6 +31,8 @@ use uuid::Uuid; #[derive(Debug)] pub struct Buck2FileCapture { child: Child, + /// Registered with the reaper for as long as the helper is ours to wait on. + _owned: crate::reaper::OwnedChild, finish_timeout: Duration, } @@ -122,8 +124,10 @@ impl Buck2FileCapture { "Invalid Buck2 capture readiness acknowledgement" )); } + let owned = crate::reaper::OwnedChild::new(child.id()); Ok(Some(Self { child, + _owned: owned, // The helper reserves five seconds to acknowledge a partial // final scan, plus five seconds for shutdown/reaping here. finish_timeout: Duration::from_secs(finalize_seconds as u64 + 10), @@ -235,8 +239,10 @@ mod tests { .kill_on_drop(true) .spawn() .err_tip(|| "Starting unresponsive test helper")?; + let owned = crate::reaper::OwnedChild::new(child.id()); let capture = Buck2FileCapture { child, + _owned: owned, finish_timeout: Duration::from_millis(20), }; let result = tokio::time::timeout(Duration::from_secs(2), capture.finish()) diff --git a/nativelink-worker/src/capacity.rs b/nativelink-worker/src/capacity.rs index 5dc7c9ebc8b..789b5f81360 100644 --- a/nativelink-worker/src/capacity.rs +++ b/nativelink-worker/src/capacity.rs @@ -1,4 +1,4 @@ -// Copyright 2024 The NativeLink Authors. All rights reserved. +// Copyright 2026 The NativeLink Authors. All rights reserved. // // Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); // you may not use this file except in compliance with the License. @@ -12,9 +12,271 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! What the worker has to spare, reported on each keepalive. +//! What this worker can see of its own CPU and memory, and what it should +//! advertise from that. +//! +//! Every deployment so far typed the advertisement by hand next to the pod's +//! limits, and the two drifted: chinchilla advertised 16 cores and 60 GiB +//! inside 15-core, 56 GiB limits. The limit is the number the kernel +//! enforces, so the advertisement is derived from it here. + +use core::hash::BuildHasher; +use std::collections::HashMap; +use std::path::{Path, PathBuf}; + +use nativelink_config::cas_server::{ + CapacityConfig, CpuUnit, MemoryEnforcement, ResourceEnforcementConfig, WorkerProperty, +}; +use nativelink_error::{Code, Error, make_err}; +use tracing::{info, warn}; + +/// CPU and memory as the cgroup (or, failing a limit, the host) reports them. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ObservedCapacity { + pub cpu_millicores: u64, + pub memory_kb: u64, +} + +/// `cpu.max` is `" "` in microseconds, or `"max "` when +/// unlimited; unlimited means the host's cores. +pub fn parse_cpu_max(cpu_max: &str, host_cpus: u64) -> Option { + let mut fields = cpu_max.split_whitespace(); + let quota = fields.next()?; + if quota == "max" { + return Some(host_cpus.saturating_mul(1000)); + } + let quota: u64 = quota.parse().ok()?; + let period: u64 = fields.next()?.parse().ok()?; + if period == 0 { + return None; + } + Some(quota.saturating_mul(1000) / period) +} + +/// `memory.max` is bytes, or `"max"` when unlimited; unlimited means the +/// host's memory. +pub fn parse_memory_max(memory_max: &str, host_memory_kb: u64) -> Option { + let value = memory_max.trim(); + if value == "max" { + return Some(host_memory_kb); + } + value.parse::().ok().map(|bytes| bytes / 1024) +} + +/// `MemTotal` from `/proc/meminfo`, in KiB. +pub fn parse_meminfo_total_kb(meminfo: &str) -> Option { + meminfo.lines().find_map(|line| { + let rest = line.strip_prefix("MemTotal:")?; + rest.split_whitespace().next()?.parse().ok() + }) +} + +/// The headroom the advertisement is divided by: the block's own figure +/// when set, else the enforcement headroom while memory enforcement is on, +/// so the two numbers cannot drift apart, else nothing. +pub fn memory_headroom_percent( + config: &CapacityConfig, + enforcement: Option<&ResourceEnforcementConfig>, +) -> u64 { + config.memory_headroom_percent.unwrap_or_else(|| { + enforcement + .filter(|enforcement| enforcement.memory == MemoryEnforcement::Soft) + .map_or(0, |enforcement| enforcement.memory_headroom_percent) + }) +} + +/// The observed capacity less the worker's own share, memory divided by +/// `memory_headroom_percent`. Returns `(cpu, memory_kb)` with the CPU on +/// the scale `cpu_unit` names: whole cores rounded down, so a 14-core pod +/// keeping one core back advertises `13`, not `13000`, to a scheduler +/// whose actions ask for `cpu_count=1`. +pub const fn advertised( + observed: ObservedCapacity, + config: &CapacityConfig, + memory_headroom_percent: u64, +) -> (u64, u64) { + let cpu_millicores = observed + .cpu_millicores + .saturating_sub(config.overhead_cpu_millicores); + let cpu = match config.cpu_unit { + CpuUnit::Cores => cpu_millicores / 1000, + CpuUnit::Millicores => cpu_millicores, + }; + let memory = observed + .memory_kb + .saturating_sub(config.overhead_memory_kb) + .saturating_mul(100) + / (100 + memory_headroom_percent); + (cpu, memory) +} + +/// What to do when the cgroup cannot be read: the configured properties +/// stand if they carry both numbers, since an operator who kept them has a +/// worker that still takes work; a worker with neither would register and +/// then satisfy no action that asks for CPU or memory, idling for good on +/// one warning line, so that is refused. +pub fn without_cgroup( + config: &CapacityConfig, + properties: &HashMap, +) -> Result<(), Error> { + let missing: Vec<&str> = [&config.cpu_property_name, &config.memory_property_name] + .into_iter() + .filter(|name| !properties.contains_key(name.as_str())) + .map(String::as_str) + .collect(); + if missing.is_empty() { + warn!( + "capacity is configured but no cgroup v2 limit applies to the worker (see the line above for the cgroup looked at); advertising the configured platform_properties instead" + ); + return Ok(()); + } + Err(make_err!( + Code::FailedPrecondition, + "capacity is configured but no cgroup v2 limit applies to the worker (not Linux, cgroup v1, no permission, or no limit at any level; see the line above), and platform_properties carries no {}: set them, or drop the capacity block", + missing.join(" or ") + )) +} + +/// The cgroup v2 path of this process from `/proc/self/cgroup`: the `0::` +/// line, which is `/` inside a container with its own cgroup namespace and +/// the full `/kubepods.slice/.../cri-containerd-.scope` path in one that +/// shares the host's, as a privileged container does. +pub fn parse_self_cgroup(self_cgroup: &str) -> Option<&str> { + self_cgroup + .lines() + .find_map(|line| line.strip_prefix("0::")) + .map(str::trim) + .filter(|path| path.starts_with('/')) +} + +/// Where this process's own `cpu.max`, `memory.max` and `memory.current` +/// live: the mount root joined with the path from `/proc/self/cgroup` when +/// that directory exists under the mount, the mount root otherwise. A +/// container with a private cgroup namespace sees itself at the root, so its +/// path is `/` and both agree. A privileged container sees the host's tree, +/// where the root's files are the host's (or absent, as on a systemd host), +/// and its own numbers sit further down. +pub fn resolve_cgroup_dir( + mount_root: &Path, + self_cgroup: Option<&str>, + is_dir: impl Fn(&Path) -> bool, +) -> PathBuf { + if let Some(path) = self_cgroup.and_then(parse_self_cgroup) { + let own = mount_root.join(path.trim_start_matches('/')); + if is_dir(&own) { + return own; + } + } + mount_root.to_path_buf() +} + +/// This process's cgroup directory on the live system. +fn cgroup_dir() -> PathBuf { + let self_cgroup = std::fs::read_to_string("/proc/self/cgroup").ok(); + resolve_cgroup_dir(Path::new("/sys/fs/cgroup"), self_cgroup.as_deref(), |dir| { + dir.join("memory.max").is_file() + }) +} + +/// Whether a `cpu.max` or `memory.max` value is "no limit": `max`, or for +/// `cpu.max` `max `. +fn is_unlimited(value: &str) -> bool { + value.split_whitespace().next() == Some("max") +} + +/// The first limit for `file` from `dir` up to and including `mount_root`, +/// with the directory it was found in. The process's own cgroup is often +/// an unlimited child of a limited parent (a systemd scope in a pod, a +/// split cgroup, a privileged container under `kubepods.slice`), and the +/// parent's limit is the one the kernel enforces on it. `None` when every +/// level says `max` or nothing is readable: no limit applies. +pub fn find_limit( + dir: &Path, + mount_root: &Path, + file: &str, + read: impl Fn(&Path) -> Option, +) -> Option<(PathBuf, String)> { + let mut at = dir; + loop { + if let Some(value) = read(&at.join(file)) + && !is_unlimited(value.trim()) + { + return Some((at.to_path_buf(), value)); + } + if at == mount_root || !at.starts_with(mount_root) { + return None; + } + at = at.parent()?; + } +} + +fn read_file(path: &Path) -> Option { + std::fs::read_to_string(path).ok() +} + +/// Reads the limits that apply to the worker: for CPU and for memory, the +/// nearest limit from its own cgroup up to the mount root. `None` where no +/// limit applies at any level (not Linux, cgroup v1, no permission, or a +/// host where nothing is limited), so the configured properties stand. +pub fn observe_cgroup() -> Option { + let host_cpus = std::thread::available_parallelism().map_or(1, |n| n.get() as u64); + let host_memory_kb = std::fs::read_to_string("/proc/meminfo") + .ok() + .and_then(|meminfo| parse_meminfo_total_kb(&meminfo)) + .unwrap_or(0); + let mount_root = Path::new("/sys/fs/cgroup"); + let dir = cgroup_dir(); + let cpu = find_limit(&dir, mount_root, "cpu.max", read_file); + let memory = find_limit(&dir, mount_root, "memory.max", read_file); + if cpu.is_none() && memory.is_none() { + warn!( + cgroup = %dir.display(), + "no cgroup v2 cpu.max or memory.max limit from this cgroup up to the mount root; keeping the configured platform_properties" + ); + return None; + } + let cpu_max = cpu.map_or_else(|| "max".to_string(), |(_, value)| value); + let memory_max = memory.map_or_else(|| "max".to_string(), |(_, value)| value); + Some(ObservedCapacity { + cpu_millicores: parse_cpu_max(&cpu_max, host_cpus)?, + memory_kb: parse_memory_max(&memory_max, host_memory_kb)?, + }) +} + +/// Sets the CPU and memory properties to what the cgroup allows. When +/// nothing can be read, the configured properties stand if they carry both, +/// and the worker fails to start if they do not. Returns what was +/// advertised. +pub fn apply( + config: &CapacityConfig, + memory_headroom_percent: u64, + properties: &mut HashMap, +) -> Result, Error> { + let Some(observed) = observe_cgroup() else { + without_cgroup(config, properties)?; + return Ok(None); + }; + let (cpu, memory_kb) = advertised(observed, config, memory_headroom_percent); + properties.insert( + config.cpu_property_name.clone(), + WorkerProperty::Values(vec![cpu.to_string()]), + ); + properties.insert( + config.memory_property_name.clone(), + WorkerProperty::Values(vec![memory_kb.to_string()]), + ); + info!( + observed_cpu_millicores = observed.cpu_millicores, + observed_memory_kb = observed.memory_kb, + cpu, + cpu_unit = ?config.cpu_unit, + memory_kb, + memory_headroom_percent, + "Advertising capacity from the cgroup" + ); + Ok(Some((cpu, memory_kb))) +} -/// `memory.current` and `memory.max` are bytes. pub fn parse_memory_current_kb(memory_current: &str) -> Option { memory_current .trim() @@ -49,17 +311,22 @@ pub const fn free_memory_kb_from( /// worker reports nothing and is never vetoed. #[cfg(target_os = "linux")] pub fn free_memory_kb() -> Option { - let limit_kb = std::fs::read_to_string("/sys/fs/cgroup/memory.max") - .ok() - .and_then(|max| { - let max = max.trim(); - (max != "max") - .then(|| parse_memory_current_kb(max)) - .flatten() - }); - let current_kb = std::fs::read_to_string("/sys/fs/cgroup/memory.current") - .ok() - .and_then(|current| parse_memory_current_kb(¤t)); + // The limit and the usage counted against it come from the same level: + // the nearest limited ancestor, which charges everything below it. + let limited = find_limit( + &cgroup_dir(), + Path::new("/sys/fs/cgroup"), + "memory.max", + read_file, + ); + let limit_kb = limited + .as_ref() + .and_then(|(_, max)| parse_memory_current_kb(max.trim())); + let current_kb = limited.as_ref().and_then(|(dir, _)| { + std::fs::read_to_string(dir.join("memory.current")) + .ok() + .and_then(|current| parse_memory_current_kb(¤t)) + }); let available_kb = std::fs::read_to_string("/proc/meminfo") .ok() .and_then(|meminfo| parse_meminfo_available_kb(&meminfo)); diff --git a/nativelink-worker/src/lib.rs b/nativelink-worker/src/lib.rs index b6a20d93847..48997ca77a8 100644 --- a/nativelink-worker/src/lib.rs +++ b/nativelink-worker/src/lib.rs @@ -20,6 +20,7 @@ pub mod local_worker; pub mod namespace_utils; pub mod persistent_worker; pub mod qos; +pub mod reaper; pub mod running_actions_manager; pub mod worker_api_client_wrapper; pub mod worker_utils; diff --git a/nativelink-worker/src/local_worker.rs b/nativelink-worker/src/local_worker.rs index c7f29b04e3a..0f93d6b86b3 100644 --- a/nativelink-worker/src/local_worker.rs +++ b/nativelink-worker/src/local_worker.rs @@ -15,13 +15,14 @@ use core::hash::BuildHasher; use core::pin::Pin; use core::str; -use core::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use core::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering}; use core::time::Duration; use std::borrow::Cow; use std::collections::HashMap; use std::env; use std::process::Stdio; use std::sync::{Arc, Weak}; +use std::time::{SystemTime, UNIX_EPOCH}; use futures::future::BoxFuture; use futures::stream::FuturesUnordered; @@ -32,27 +33,30 @@ use nativelink_metric::{MetricsComponent, RootMetricsComponent}; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::update_for_worker::Update; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::worker_api_client::WorkerApiClient; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{ - ActionResourceUsage, ExecuteComplete, ExecuteResult, GoingAwayRequest, KeepAliveRequest, - UpdateForWorker, WorkerLoad, execute_result, + ActionResourceUsage, ExecuteAccepted, ExecuteComplete, ExecuteDeclined, ExecuteResult, + GoingAwayRequest, KeepAliveRequest, StartExecute, UpdateForWorker, WorkerLoad, + execute_declined, execute_result, }; use nativelink_store::fast_slow_store::FastSlowStore; use nativelink_util::action_messages::{ActionResult, ActionStage, OperationId}; use nativelink_util::common::fs; use nativelink_util::digest_hasher::DigestHasherFunc; +use nativelink_util::health_utils::{HealthStatus, HealthStatusIndicator}; use nativelink_util::metrics_utils::{AsyncCounterWrapper, CounterWithTime}; use nativelink_util::shutdown_guard::ShutdownGuard; use nativelink_util::store_trait::Store; -use nativelink_util::{spawn, tls_utils}; +use nativelink_util::{background_spawn, spawn, spawn_blocking, tls_utils}; use opentelemetry::context::Context; use tokio::sync::{broadcast, mpsc}; use tokio::{process, time}; use tokio_stream::wrappers::UnboundedReceiverStream; -use tonic::Streaming; +use tonic::{Streaming, async_trait}; use tracing::{Level, debug, error, event, info, info_span, instrument, trace, warn}; use crate::running_actions_manager::{ - ExecutionConfiguration, Metrics as RunningActionManagerMetrics, RunningAction, - RunningActionsManager, RunningActionsManagerArgs, RunningActionsManagerImpl, + ExecutionConfiguration, MISSING_INPUT_ERROR_TIP, Metrics as RunningActionManagerMetrics, + ResourceEnforcement, RunningAction, RunningActionsManager, RunningActionsManagerArgs, + RunningActionsManagerImpl, }; use crate::worker_api_client_wrapper::{WorkerApiClientTrait, WorkerApiClientWrapper}; use crate::worker_utils::make_connect_worker_request; @@ -84,8 +88,23 @@ impl Drop for ActionsInTransitGuard { } /// If we lose connection to the worker api server we will wait this many seconds -/// before trying to connect. +/// before trying to connect, doubling on every failed attempt up to +/// `CONNECTION_RETRY_MAX_DELAY_S`, with jitter so a fleet that lost its +/// scheduler together does not redial together. const CONNECTION_RETRY_DELAY_S: f32 = 0.5; +const CONNECTION_RETRY_MAX_DELAY_S: f32 = 30.0; + +/// Delay before the `attempt`th consecutive reconnect (0 = first retry), +/// between half and one and a half times the exponential figure. +fn reconnect_delay(attempt: u32) -> Duration { + let exponent = i32::try_from(attempt.min(16)).unwrap_or(16); + let base = (CONNECTION_RETRY_DELAY_S * 2f32.powi(exponent)).min(CONNECTION_RETRY_MAX_DELAY_S); + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |d| d.subsec_nanos()); + let jitter = 0.5 + (nanos % 1_000) as f32 / 1_000.0; + Duration::from_secs_f32(base * jitter) +} /// Default endpoint timeout. If this value gets modified the documentation in /// `cas_server.rs` must also be updated. @@ -95,8 +114,14 @@ const DEFAULT_ENDPOINT_TIMEOUT_S: f32 = 5.; /// If this value gets modified the documentation in `cas_server.rs` must also be updated. const DEFAULT_MAX_ACTION_TIMEOUT: Duration = Duration::from_mins(20); const DEFAULT_MAX_UPLOAD_TIMEOUT: Duration = Duration::from_mins(10); +/// Default for `max_download_timeout_s`: an input fetch is quick or wedged. +const DEFAULT_MAX_DOWNLOAD_TIMEOUT: Duration = Duration::from_mins(10); const DEFAULT_MAX_CLEANUP_WAIT: Duration = Duration::from_secs(30); const DEFAULT_MAX_CLEANUP_BACKOFF: Duration = Duration::from_millis(500); +/// If this value gets modified the documentation in `cas_server.rs` must also be updated. +const DEFAULT_PRECONDITION_TIMEOUT: Duration = Duration::from_secs(30); +/// If this value gets modified the documentation in `cas_server.rs` must also be updated. +const DEFAULT_KILL_GRACE: Duration = Duration::from_secs(5); struct FinishedActionResult { action_result: ActionResult, @@ -115,12 +140,78 @@ struct LocalWorkerImpl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsM // on by the scheduler. actions_in_transit: Arc, accepted_action: AtomicBool, + /// The scheduler said it understands `ExecuteAccepted` and + /// `ExecuteDeclined`. Without it the worker runs whatever it is sent, + /// as every earlier release did. + dispatch_ack: bool, + /// The platform property the scheduler reads an action's memory + /// reservation from, as it told us on connection; empty when it does + /// not veto on memory. Read from the same property, a refusal for load + /// here agrees with what the scheduler would have vetoed. + memory_property: String, metrics: Arc, } +/// Why this worker will not run an action it was just sent. +enum Refusal { + AtCapacity { in_flight: u64, max: u64 }, + Load { needed_kb: u64, free_kb: u64 }, + ShuttingDown, +} + +impl Refusal { + fn into_declined(self, operation_id: String) -> ExecuteDeclined { + let (reason, detail, needed_kb, free_kb) = match self { + Self::AtCapacity { in_flight, max } => ( + execute_declined::Reason::AtCapacity, + format!("{in_flight} of {max} in flight"), + 0, + 0, + ), + Self::Load { needed_kb, free_kb } => ( + execute_declined::Reason::Load, + format!("needs {needed_kb} KiB, {free_kb} KiB free"), + needed_kb, + free_kb, + ), + Self::ShuttingDown => ( + execute_declined::Reason::ShuttingDown, + "worker shutting down".to_string(), + 0, + 0, + ), + }; + ExecuteDeclined { + operation_id, + reason: reason as i32, + detail, + needed_kb, + free_kb, + } + } +} + +/// The memory reservation an action carries under `property`, in KiB, if +/// it declares one. An empty property name is a scheduler that does not +/// veto on memory, so nothing is ever read. +fn memory_reservation_kb(start_execute: &StartExecute, property: &str) -> Option { + if property.is_empty() { + return None; + } + start_execute + .platform + .as_ref()? + .properties + .iter() + .find(|p| p.name == property) + .and_then(|p| p.value.parse::().ok()) + .filter(|kb| *kb > 0) +} + pub async fn preconditions_met( precondition_script: Option, extra_envs: &HashMap, + timeout: Duration, ) -> Result<(), Error> { let Some(precondition_script) = &precondition_script else { // No script means we are always ok to proceed. @@ -154,7 +245,19 @@ pub async fn preconditions_met( .envs(extra_envs) .spawn() .err_tip(|| format!("Could not execute precondition command {precondition_script:?}"))?; - let output = precondition_process.wait_with_output().await?; + let _owned = crate::reaper::OwnedChild::new(precondition_process.id()); + // Bounded: a script that hangs held the action forever, and + // `kill_on_drop` ends the script when the timeout drops it. + let output = match time::timeout(timeout, precondition_process.wait_with_output()).await { + Ok(output) => output?, + Err(_) => { + return Err(make_err!( + Code::ResourceExhausted, + "Preconditions script {precondition_script:?} did not finish within {} ms", + timeout.as_millis() + )); + } + }; let stdout = str::from_utf8(&output.stdout).unwrap_or(""); trace!(status = %output.status, %stdout, "Preconditions script returned"); if output.status.code() == Some(0) { @@ -174,6 +277,8 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke config: &'a LocalWorkerConfig, grpc_client: T, worker_id: String, + dispatch_ack: bool, + memory_property: String, running_actions_manager: Arc, metrics: Arc, ) -> Self { @@ -181,6 +286,8 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke config, grpc_client, worker_id, + dispatch_ack, + memory_property, running_actions_manager, // Number of actions that have been received in `Update::StartAction`, but // not yet processed by running_actions_manager's spawn. This number should @@ -192,6 +299,35 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke } } + /// Local admission: the worker's own word on whether it can take this + /// action now. The scheduler's ledger says what it believes the worker + /// has; this is what the worker has. + fn admission(&self, start_execute: &StartExecute, in_flight: u64) -> Option { + let max = self.config.max_inflight_tasks; + if max > 0 && in_flight >= max { + return Some(Refusal::AtCapacity { in_flight, max }); + } + if let (Some(needed_kb), Some(free_kb)) = ( + memory_reservation_kb(start_execute, &self.memory_property), + crate::capacity::free_memory_kb(), + ) && free_kb < needed_kb + { + return Some(Refusal::Load { needed_kb, free_kb }); + } + None + } + + /// Tells the scheduler the action is refused; only meaningful when it + /// understands the message. + async fn decline(&self, operation_id: String, refusal: Refusal) -> Result<(), Error> { + self.metrics.actions_declined.inc(); + self.grpc_client + .clone() + .execute_declined(refusal.into_declined(operation_id)) + .await + .err_tip(|| "Could not send ExecuteDeclined") + } + /// Starts a background spawn/thread that will send a message to the server every `timeout / 2`. async fn start_keep_alive(&self) -> Result<(), Error> { // According to tonic's documentation this call should be cheap and is the same stream. @@ -279,7 +415,7 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke if let Err(err) = self .grpc_client .clone() - .going_away(GoingAwayRequest {}) + .going_away(GoingAwayRequest { drain: false }) .await { warn!(?err, "Could not unregister completed single-use worker"); @@ -319,8 +455,10 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke Update::StartAction(start_execute) => { // Don't accept any new requests if we're shutting down. if shutting_down || (self.config.single_use - && self.accepted_action.swap(true, Ordering::AcqRel)) { - if let Some(instance_name) = start_execute.execute_request.map(|request| request.instance_name) { + && self.accepted_action.load(Ordering::Acquire)) { + if self.dispatch_ack { + self.decline(start_execute.operation_id, Refusal::ShuttingDown).await?; + } else if let Some(instance_name) = start_execute.execute_request.map(|request| request.instance_name) { self.grpc_client.clone().execution_response( ExecuteResult{ instance_name, @@ -333,6 +471,34 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke continue; } + // Admission, then the acknowledgement: what the + // scheduler charged on the send is confirmed or + // handed back before anything runs. A scheduler + // that does not speak the acknowledgement gets + // the old behaviour, run whatever arrives. + if self.dispatch_ack { + if let Some(refusal) = self.admission( + &start_execute, + actions_in_flight.load(Ordering::Acquire), + ) { + self.decline(start_execute.operation_id, refusal).await?; + continue; + } + self.grpc_client + .clone() + .execute_accepted(ExecuteAccepted { + operation_id: start_execute.operation_id.clone(), + }) + .await + .err_tip(|| "Could not send ExecuteAccepted")?; + } + // Admitted: a single-use worker is spent from + // here. A decline above must not spend it, or + // every dispatch it turns away costs a pod. + if self.config.single_use { + self.accepted_action.store(true, Ordering::Release); + } + self.metrics.start_actions_received.inc(); let execute_request = start_execute.execute_request.as_ref(); @@ -347,6 +513,11 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke let start_action_fut = { let precondition_script_cfg = self.config.experimental_precondition_script.clone(); + let precondition_timeout = if self.config.precondition_timeout_ms == 0 { + DEFAULT_PRECONDITION_TIMEOUT + } else { + Duration::from_millis(self.config.precondition_timeout_ms) + }; let mut extra_envs: HashMap = HashMap::new(); if let Some(ref additional_environment) = self.config.additional_environment { for (name, source) in additional_environment { @@ -374,7 +545,7 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke }; let single_use = self.config.single_use; self.metrics.clone().wrap(move |metrics| async move { - metrics.preconditions.wrap(preconditions_met(precondition_script_cfg, &extra_envs)) + metrics.preconditions.wrap(preconditions_met(precondition_script_cfg, &extra_envs, precondition_timeout)) .and_then(|()| running_actions_manager.create_and_add_action(worker_id, start_execute)) .map(move |r| { // Now that we either failed or registered our action, we can @@ -456,8 +627,10 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke .err_tip(|| "Error while calling execution_response")?; }, Err(e) => { + // The manager tags a NotFound on anything the client + // referenced by digest; no store's wording involved. let is_cas_blob_missing = e.code == Code::NotFound - && e.message_string().contains("not found in either fast or slow store"); + && e.messages.iter().any(|m| m == MISSING_INPUT_ERROR_TIP); if is_cas_blob_missing { warn!( ?e, @@ -549,16 +722,42 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke Error::from_std_err(Code::Internal, &e).append("Failed to receive shutdown message"))?; let actions_in_flight = actions_in_flight.clone(); let actions_notify = actions_notify.clone(); + let drain_on_shutdown = self.config.drain_on_shutdown; + let drain_deadline = if self.config.max_action_timeout_s == 0 { + DEFAULT_MAX_ACTION_TIMEOUT + } else { + Duration::from_secs(self.config.max_action_timeout_s as u64) + }; let shutdown_future = async move { - // Wait for in-flight operations to be fully completed. - while actions_in_flight.load(Ordering::Acquire) > 0 { - actions_notify.notified().await; + if drain_on_shutdown { + // Say so first, so nothing new is dispatched here + // while the running actions finish; the scheduler + // removes this worker when the stream closes. + if let Err(e) = grpc_client.going_away(GoingAwayRequest { drain: true }).await { + error!("Failed to send GoingAwayRequest: {e}",); + return Err(e); + } + } + // Wait for in-flight operations to be fully completed, + // for as long as one action is allowed to run. + let wait = async { + while actions_in_flight.load(Ordering::Acquire) > 0 { + actions_notify.notified().await; + } + }; + if time::timeout(drain_deadline, wait).await.is_err() { + error!( + actions_in_flight = actions_in_flight.load(Ordering::Acquire), + "Drain deadline passed with actions still in flight; shutting down anyway" + ); } - // Sending this message immediately evicts all jobs from - // this worker, of which there should be none. - if let Err(e) = grpc_client.going_away(GoingAwayRequest {}).await { - error!("Failed to send GoingAwayRequest: {e}",); - return Err(e); + if !drain_on_shutdown { + // Sending this message immediately evicts all jobs from + // this worker, of which there should be none. + if let Err(e) = grpc_client.going_away(GoingAwayRequest { drain: false }).await { + error!("Failed to send GoingAwayRequest: {e}",); + return Err(e); + } } // Allow shutdown to occur now. drop(shutdown_guard); @@ -581,6 +780,80 @@ pub struct LocalWorker, sleep_fn: Option BoxFuture<'static, ()> + Send + Sync>>, metrics: Arc, + registration: Arc, +} + +/// Whether the worker currently holds a registration with its scheduler. +/// A worker that has not registered, or lost its connection and is +/// reconnecting, cannot take work; as a health indicator this is what a +/// readiness probe reads. +#[derive(Debug)] +pub struct WorkerRegistration { + name: String, + registered: AtomicBool, + ever_registered: AtomicBool, +} + +impl WorkerRegistration { + pub fn new(name: &str) -> Arc { + Arc::new(Self { + name: name.to_string(), + registered: AtomicBool::new(false), + ever_registered: AtomicBool::new(false), + }) + } + + pub fn is_registered(&self) -> bool { + self.registered.load(Ordering::Acquire) + } + + fn set_registered(&self, registered: bool) { + self.registered.store(registered, Ordering::Release); + if registered { + self.ever_registered.store(true, Ordering::Release); + } + } +} + +#[async_trait] +impl HealthStatusIndicator for WorkerRegistration { + fn get_name(&self) -> &'static str { + "WorkerRegistration" + } + + async fn check_health(&self, _namespace: Cow<'static, str>) -> HealthStatus { + if self.is_registered() { + HealthStatus::Ok { + struct_name: "WorkerRegistration", + message: Cow::Owned(format!( + "worker '{}' registered with the scheduler", + self.name + )), + } + } else if self.ever_registered.load(Ordering::Acquire) { + // Lost after it was there: the worker is reconnecting, and until + // it does it holds no work. Initializing, not Failed: Failed + // turns the plain status check red too, and a liveness probe on + // it would restart every worker during a scheduler roll longer + // than its threshold, and redden a co-hosted CAS on one worker's + // blip. This belongs to readiness alone. + HealthStatus::Initializing { + struct_name: "WorkerRegistration", + message: Cow::Owned(format!( + "worker '{}' lost its scheduler connection, reconnecting", + self.name + )), + } + } else { + HealthStatus::Initializing { + struct_name: "WorkerRegistration", + message: Cow::Owned(format!( + "worker '{}' not yet registered with the scheduler", + self.name + )), + } + } + } } impl< @@ -653,6 +926,11 @@ pub async fn new_local_worker( } else { Duration::from_secs(config.max_upload_timeout_s as u64) }; + let max_download_timeout = if config.max_download_timeout_s == 0 { + DEFAULT_MAX_DOWNLOAD_TIMEOUT + } else { + Duration::from_secs(config.max_download_timeout_s as u64) + }; let max_cleanup_wait = if config.max_cleanup_wait_s == 0 { DEFAULT_MAX_CLEANUP_WAIT } else { @@ -705,31 +983,50 @@ pub async fn new_local_worker( }; #[cfg(target_os = "linux")] - let use_namespaces = if let Some(use_namespaces) = &config.use_namespaces { - if *use_namespaces - && !crate::namespace_utils::namespaces_supported( - config.use_mount_namespace.unwrap_or_default(), - ) - { - return Err(make_err!(Code::Unavailable, "Namespaces not supported")); + let use_namespaces = { + let use_mount_namespace = config.use_mount_namespace.unwrap_or_default(); + // The mount namespace only exists when `use_namespaces` is on too; + // `isolate_tmp` is judged against that, not against the flag alone, + // or an explicit `true` would be dropped without a word. + let mount_namespace_on = config.use_namespaces == Some(true) && use_mount_namespace; + // A private /tmp is part of the mount isolation unless explicitly + // turned off. A worker without a /tmp, such as a minimal container + // image, has nothing for actions to collide on, so the default is + // off there instead of failing to start. + let isolate_tmp = config.isolate_tmp.unwrap_or_else(|| { + let has_tmp = std::path::Path::new("/tmp").is_dir(); + if mount_namespace_on && !has_tmp { + warn!("/tmp does not exist on this worker, so actions will not get a private /tmp"); + } + mount_namespace_on && has_tmp + }); + if isolate_tmp && !mount_namespace_on { + return Err(make_err!( + Code::InvalidArgument, + "isolate_tmp requires use_namespaces and use_mount_namespace to be true" + )); } - if !*use_namespaces { - crate::running_actions_manager::UseNamespaces::No - } else if config.use_mount_namespace.unwrap_or_default() { - crate::running_actions_manager::UseNamespaces::YesAndMount + if let Some(use_namespaces) = &config.use_namespaces { + if *use_namespaces + && !crate::namespace_utils::namespaces_supported(use_mount_namespace, isolate_tmp) + { + return Err(make_err!(Code::Unavailable, "Namespaces not supported")); + } + if !*use_namespaces { + crate::running_actions_manager::UseNamespaces::No + } else if use_mount_namespace { + crate::running_actions_manager::UseNamespaces::YesAndMount { isolate_tmp } + } else { + crate::running_actions_manager::UseNamespaces::Yes + } + } else if use_mount_namespace { + return Err(make_err!( + Code::Unavailable, + "Mount namespaces not supported" + )); } else { - crate::running_actions_manager::UseNamespaces::Yes + crate::running_actions_manager::UseNamespaces::No } - } else if config - .use_mount_namespace - .is_some_and(core::convert::identity) - { - return Err(make_err!( - Code::Unavailable, - "Mount namespaces not supported" - )); - } else { - crate::running_actions_manager::UseNamespaces::No }; #[cfg(not(target_os = "linux"))] @@ -749,14 +1046,46 @@ pub async fn new_local_worker( "Mount namespaces not supported on non-Linux OSes" )); } + #[cfg(not(target_os = "linux"))] + if config.isolate_tmp.is_some_and(core::convert::identity) { + return Err(make_err!( + Code::Unavailable, + "isolate_tmp is not supported on non-Linux OSes" + )); + } + + // A pooled worker process gets the same PID, user, UTS and IPC + // namespaces a one-shot action gets, never the mount namespace. + #[cfg(target_os = "linux")] + let persistent_workers_namespaced = !matches!( + use_namespaces, + crate::running_actions_manager::UseNamespaces::No + ); + #[cfg(not(target_os = "linux"))] + let persistent_workers_namespaced = false; let running_actions_manager = Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { root_action_directory: config.work_directory.clone(), execution_configuration: ExecutionConfiguration { + max_captured_output_bytes: config.max_captured_output_bytes, + kill_grace: if config.kill_grace_ms == 0 { + DEFAULT_KILL_GRACE + } else { + Duration::from_millis(config.kill_grace_ms) + }, + set_tmpdir: config.set_tmpdir, + persistent_workers: persistent_worker_settings( + config.persistent_workers.as_ref(), + persistent_workers_namespaced, + ), buck2_file_capture: config.experimental_buck2_file_capture.clone(), entrypoint, additional_environment: config.additional_environment.clone(), + resource_enforcement: config + .resource_enforcement + .as_ref() + .and_then(ResourceEnforcement::from_config), }, cas_store: fast_slow_store, ac_store, @@ -764,6 +1093,7 @@ pub async fn new_local_worker( upload_action_result_config: &config.upload_action_result, max_action_timeout, max_upload_timeout, + max_download_timeout, max_cleanup_wait, max_cleanup_backoff, timeout_handled_externally: config.timeout_handled_externally, @@ -772,6 +1102,46 @@ pub async fn new_local_worker( #[cfg(target_os = "linux")] use_namespaces, })?); + // What actions leave behind is reaped here; see `crate::reaper`. + if let Err(err) = crate::reaper::become_subreaper() { + warn!( + ?err, + "Could not become the subreaper; orphans of actions will not be reaped here" + ); + } + drop(background_spawn!("orphan_reaper", async move { + let mut seen_last = std::collections::BTreeSet::new(); + loop { + time::sleep(crate::reaper::REAP_INTERVAL).await; + // The sweep reads /proc for every process; off the runtime. + let (reaped, seen) = spawn_blocking!("orphan_reaper_sweep", move || { + crate::reaper::reap_orphaned_zombies(&seen_last) + }) + .await + .unwrap_or_default(); + seen_last = seen; + if reaped > 0 { + info!(reaped, "Reaped zombies that actions left behind"); + } + } + })); + if config.orphan_sweep_interval_s > 0 { + let interval = Duration::from_secs(config.orphan_sweep_interval_s); + let manager = running_actions_manager.clone(); + // Detached on purpose: the guarded spawn aborts its task when the + // handle drops, and this one runs for the worker's whole life. + drop(background_spawn!("orphan_sweep", async move { + info!(interval_s = interval.as_secs(), "Orphan sweep scheduled"); + loop { + time::sleep(interval).await; + match manager.sweep_orphaned_directories().await { + Ok(0) => debug!("Orphan sweep found nothing"), + Ok(removed) => info!(removed, "Orphan sweep finished"), + Err(err) => warn!(?err, "Orphan sweep failed"), + } + } + })); + } let local_worker = LocalWorker::new_with_connection_factory_and_actions_manager( config.clone(), running_actions_manager, @@ -819,15 +1189,30 @@ impl LocalWorker Arc { + self.registration.clone() + } + + /// Flip a flag that was registered with a health registry before the + /// worker existed, as the server binary has to. + #[must_use] + pub fn with_registration(mut self, registration: Arc) -> Self { + self.registration = registration; + self + } + #[allow( clippy::missing_const_for_fn, reason = "False positive on stable, but not on nightly" @@ -839,7 +1224,7 @@ impl LocalWorker Result<(String, Streaming), Error> { + ) -> Result<(String, bool, String, Streaming), Error> { let mut extra_envs: HashMap = HashMap::new(); if let Some(ref additional_environment) = self.config.additional_environment { for (name, source) in additional_environment { @@ -860,9 +1245,18 @@ impl LocalWorker LocalWorker connection_result.worker_id, + let (worker_id, dispatch_ack, memory_property) = match first_msg_update { + Some(Update::ConnectionResult(connection_result)) => ( + connection_result.worker_id, + connection_result.dispatch_ack, + connection_result.memory_property, + ), other => { return Err(make_input_err!( "Expected first response from scheduler to be a ConnectionResult got : {:?}", @@ -893,7 +1291,12 @@ impl LocalWorker LocalWorker LocalWorker ( + Ok((worker_id, dispatch_ack, memory_property, update_for_worker_stream)) => ( LocalWorkerImpl::new( &self.config, client, worker_id, + dispatch_ack, + memory_property, self.running_actions_manager.clone(), self.metrics.clone(), ), @@ -952,9 +1361,13 @@ impl LocalWorker, + namespaced: bool, +) -> crate::running_actions_manager::PersistentWorkersSettings { + use crate::persistent_worker::PoolConfig; + let defaults = PoolConfig::default(); + let pool = config.map_or(defaults, |c| { + let or_default = |value: u64, default: u64| if value == 0 { default } else { value }; + PoolConfig { + max_workers_per_key: if c.max_workers_per_key == 0 { + defaults.max_workers_per_key + } else { + c.max_workers_per_key + }, + idle_timeout: Duration::from_secs(or_default( + c.idle_timeout_s, + defaults.idle_timeout.as_secs(), + )), + max_requests_per_worker: or_default( + c.max_requests_per_worker, + defaults.max_requests_per_worker, + ), + shutdown_grace: Duration::from_millis(or_default( + c.shutdown_grace_ms, + u64::try_from(defaults.shutdown_grace.as_millis()).unwrap_or(5000), + )), + acquire_timeout: Duration::from_secs(or_default( + c.acquire_timeout_s, + defaults.acquire_timeout.as_secs(), + )), + namespaced, + } + }); + crate::running_actions_manager::PersistentWorkersSettings { + enabled: config.is_none_or(|c| c.enabled), + pool: PoolConfig { namespaced, ..pool }, + } +} diff --git a/nativelink-worker/src/namespace_utils.rs b/nativelink-worker/src/namespace_utils.rs index bb9d34208b7..206759a92ee 100644 --- a/nativelink-worker/src/namespace_utils.rs +++ b/nativelink-worker/src/namespace_utils.rs @@ -16,14 +16,38 @@ use std::io::Error; use tracing::error; -/// A wrapper around a Child to send SIGTERM to kill the process instead -/// of SIGKILL as it's wrapped by the stub. +/// A wrapper around a Child that may be the stub of a PID namespace. A +/// SIGTERM to the stub is passed on to the action; a SIGKILL to the stub +/// ends the whole namespace, which is the hard kill. #[derive(Debug)] pub struct MaybeNamespacedChild { namespaced: bool, child: tokio::process::Child, } +/// SIGKILL to every process in the group `pgid` leads. ESRCH means the +/// group is already gone, which is the outcome wanted; anything else, EPERM +/// above all, means a survivor and is reported. +#[cfg(target_os = "linux")] +fn kill_process_group(pgid: u32) { + let Ok(pgid) = i32::try_from(pgid) else { + return; + }; + // SAFETY: killpg only takes integers and has no memory safety + // considerations; a stale group id is reported as ESRCH, not acted on. + let rc = unsafe { libc::killpg(pgid, libc::SIGKILL) }; + if rc != 0 { + let err = Error::last_os_error(); + if err.raw_os_error() != Some(libc::ESRCH) { + error!( + pgid, + ?err, + "Could not kill the action's process group; a descendant may have survived" + ); + } + } +} + impl MaybeNamespacedChild { pub const fn new(namespaced: bool, child: tokio::process::Child) -> Self { Self { namespaced, child } @@ -33,27 +57,19 @@ impl MaybeNamespacedChild { self.child.id() } - /// Send SIGTERM if namespaced which sends SIGKILL to the child, otherwise - /// send SIGKILL to the child. + /// SIGKILL the child and wait for it. Namespaced, the child is the + /// stub and the init of the action's PID namespace, so its death takes + /// the action and everything the action spawned with it; the stub also + /// set `PR_SET_PDEATHSIG` on the action for the same end. Not + /// namespaced, the child is its own group leader, so the group gets the + /// SIGKILL first: a timed-out or cancelled action used to lose only its + /// direct child while its descendants ran on (issue #225). pub async fn kill(&mut self) -> Result<(), Error> { - if self.namespaced { - // It would be safer to call send_signal to use the pidfd to avoid - // races, however this is still an experimental API, see: - // https://github.com/rust-lang/rust/issues/141975 - // self.child.std_child().send_signal(Signal::SIGTERM)?; - // return self.child.wait().await.map(|_| ()); - if let Some(pid) = self.child.id() { - let pid_t: libc::pid_t = pid.try_into().map_err(|e| { - Error::new( - std::io::ErrorKind::InvalidInput, - format!("pid larger than pid_t type ({pid}): {e}"), - ) - })?; - // SAFETY: pid is valid as provided by the wrapper and we are - // sending a signal to the namespaced stub. - unsafe { libc::kill(pid_t, libc::SIGTERM) }; - return self.child.wait().await.map(|_| ()); - } + #[cfg(target_os = "linux")] + if !self.namespaced + && let Some(pgid) = self.child.id() + { + kill_process_group(pgid); } self.child.kill().await } @@ -83,11 +99,26 @@ const NS_ERROR_TYPE_MASK: i32 = 0x3; // 11 - i.e. NS_ERROR_TYPE_BITS lowest bits /// Determines whether the namespaces provided by this module are supported /// on the currently running system by forking a process and trying to enter -/// it into the new namespaces. -pub fn namespaces_supported(mount: bool) -> bool { - // SAFETY: Posix requires that geteuid is always successful. +/// it into the new namespaces. When `mount` is set the mount namespace is +/// checked as well, and when `isolate_tmp` is also set the check includes +/// bind mounting over `/tmp`. +pub fn namespaces_supported(mount: bool, isolate_tmp: bool) -> bool { + if isolate_tmp && !mount { + error!("Namespaces: isolating /tmp requires a mount namespace"); + return false; + } + if isolate_tmp && !std::path::Path::new("/tmp").is_dir() { + error!("Namespaces: isolating /tmp requires /tmp to exist and be a directory"); + return false; + } + // SAFETY: Posix requires that geteuid and getegid are always successful. let uid = unsafe { libc::geteuid() }; + // SAFETY: As above. + let gid = unsafe { libc::getegid() }; + // Read before the fork: once the child has unshared, its ids are the + // unmapped overflow ids until the maps are written. let uid_map = format!("{uid} {uid} 1\n"); + let gid_map = format!("{gid} {gid} 1\n"); // SAFETY: We ensure that if pid == 0 we only call async-signal-safe functions. let pid = unsafe { libc::fork() }; match pid { @@ -100,8 +131,32 @@ pub fn namespaces_supported(mount: bool) -> bool { // SAFETY: Unshare does not have any unsafe effects and modifies no // memory, it is also async-signal-safe. if unsafe { libc::unshare(flags) } == 0 { + // The maps `configure_namespace` writes for an action, in + // its order and with its tolerance. The gid has to be mapped + // too: an inode created on a filesystem mounted inside the + // namespace, as the mkdir below does, needs the creator's + // group to map, or the kernel refuses it with EOVERFLOW. + if let Err(err) = write_signal_safe(c"/proc/self/setgroups", b"deny") + && err != libc::EPERM + && err != libc::EACCES + && err != libc::ENOENT + { + exit( + (NamespaceErrorType::WriteSignalSafe as i32) | (err << NS_ERROR_TYPE_BITS), + ); + } match write_signal_safe(c"/proc/self/uid_map", uid_map.as_bytes()) { Ok(()) => { + if let Err(err) = + write_signal_safe(c"/proc/self/gid_map", gid_map.as_bytes()) + && err != libc::EPERM + && err != libc::EACCES + { + exit( + (NamespaceErrorType::WriteSignalSafe as i32) + | (err << NS_ERROR_TYPE_BITS), + ); + } if !mount { exit(0); } @@ -114,13 +169,32 @@ pub fn namespaces_supported(mount: bool) -> bool { libc::MS_REC | libc::MS_PRIVATE, core::ptr::null(), ) - } == 0 + } != 0 { - exit(0); + // SAFETY: We just called a libc function that failed (-1). + let errno = unsafe { *libc::__errno_location() }; + exit( + (NamespaceErrorType::Mount as i32) | (errno << NS_ERROR_TYPE_BITS), + ); } - // SAFETY: We just called a libc function that failed (-1). - let errno = unsafe { *libc::__errno_location() }; - exit((NamespaceErrorType::Mount as i32) | (errno << NS_ERROR_TYPE_BITS)); + // The same sequence `perform_remount` runs for an + // action, so a policy that allows a bind to self but + // refuses a sized tmpfs, or a mkdir under it, fails + // here and not in every action's pre_exec. All of it + // is inside this child's private mount namespace. + if isolate_tmp + && let Err(err) = bind_mount(c"/tmp", c"/tmp") + .and_then(|()| mount_tmpfs(c"/tmp", Some(ROOT_MASK_TMPFS_OPTIONS))) + .and_then(|()| { + mkdir_p_signal_safe(c"/tmp/nativelink-namespace-probe/action") + }) + { + let errno = err.raw_os_error().unwrap_or(libc::EIO); + exit( + (NamespaceErrorType::Mount as i32) | (errno << NS_ERROR_TYPE_BITS), + ); + } + exit(0); } Err(uid_map_err) => { exit( @@ -162,7 +236,7 @@ pub fn namespaces_supported(mount: bool) -> bool { s if s & NS_ERROR_TYPE_MASK == NamespaceErrorType::WriteSignalSafe as i32 => { error!( errno = s >> NS_ERROR_TYPE_BITS, - "Namespaces: Error while writing to /proc/self/uid_map" + "Namespaces: Error while writing the id maps under /proc/self" ); } s if s & NS_ERROR_TYPE_MASK == NamespaceErrorType::Mount as i32 => { @@ -298,7 +372,90 @@ impl Drop for OwnedFd { } } +/// With a private `/tmp`, the tmpfs that masks the root action directory +/// holds one directory entry and a mount point, so it is bounded to this: +/// anything that writes beside the action directory instead of inside it +/// gets `ENOSPC` rather than the worker's memory. Without `isolate_tmp` +/// the mask is unbounded, as it was before the private `/tmp` existed. +const ROOT_MASK_TMPFS_OPTIONS: &core::ffi::CStr = c"size=1m"; + +/// Whether `path` is `/tmp` or lies under it. +fn is_under_tmp(path: &core::ffi::CStr) -> bool { + let bytes = path.to_bytes(); + bytes == b"/tmp" || bytes.starts_with(b"/tmp/") +} + +/// Mount a fresh, empty tmpfs over the given directory in an +/// async-signal-safe manner, with the given mount options, or none. +fn mount_tmpfs(target: &core::ffi::CStr, options: Option<&core::ffi::CStr>) -> Result<(), Error> { + // SAFETY: mount is async-signal-safe. The filesystem type, target and options are valid C-strings. + if unsafe { + libc::mount( + c"tmpfs".as_ptr(), + target.as_ptr(), + c"tmpfs".as_ptr(), + 0, + options.map_or(core::ptr::null(), |options| options.as_ptr().cast()), + ) + } != 0 + { + return Err(Error::last_os_error()); + } + Ok(()) +} + +/// Bind mount `source` over `target` in an async-signal-safe manner. +fn bind_mount(source: &core::ffi::CStr, target: &core::ffi::CStr) -> Result<(), Error> { + // SAFETY: mount is async-signal-safe. Both paths are valid C-strings. + if unsafe { + libc::mount( + source.as_ptr(), + target.as_ptr(), + core::ptr::null(), + libc::MS_BIND | libc::MS_REC, + core::ptr::null(), + ) + } != 0 + { + return Err(Error::last_os_error()); + } + Ok(()) +} + +/// Create a directory and any missing parents in an async-signal-safe +/// manner, like `mkdir -p`. Components that already exist are skipped, so +/// this is a no-op for a path that is already present. +fn mkdir_p_signal_safe(path: &core::ffi::CStr) -> Result<(), Error> { + let bytes = path.to_bytes(); + // Leave room for the NUL that terminates each prefix. + let mut buffer = [0u8; libc::PATH_MAX as usize]; + if bytes.len() >= buffer.len() { + return Err(Error::from_raw_os_error(libc::ENAMETOOLONG)); + } + buffer[..bytes.len()].copy_from_slice(bytes); + // Start at 1 so a leading '/' is never treated as an empty component. + for i in 1..=bytes.len() { + if i < bytes.len() && buffer[i] != b'/' { + continue; + } + // Temporarily terminate the path here so just this prefix is created. + buffer[i] = 0; + // SAFETY: mkdir is async-signal-safe and the buffer is NUL-terminated at i. + if unsafe { libc::mkdir(buffer.as_ptr().cast(), 0o777) } != 0 { + let err = Error::last_os_error(); + if err.raw_os_error() != Some(libc::EEXIST) { + return Err(err); + } + } + if i < bytes.len() { + buffer[i] = b'/'; + } + } + Ok(()) +} + fn perform_remount( + tmp_directory: Option<&core::ffi::CStr>, root_action_directory: &core::ffi::CStr, action_directory: &core::ffi::CStr, ) -> Result<(), Error> { @@ -342,19 +499,37 @@ fn perform_remount( } let fd = OwnedFd(fd); - // Mask the root action directory with a tmpfs to ensure sibling directories aren't visible. - // SAFETY: mount is async-signal-safe. The filesystem type and target are valid C-strings. - if unsafe { - libc::mount( - c"tmpfs".as_ptr(), - root_action_directory.as_ptr(), - c"tmpfs".as_ptr(), - 0, - core::ptr::null(), - ) - } != 0 - { - return Err(Error::last_os_error()); + if let Some(tmp_directory) = tmp_directory { + // Give the action a private /tmp: its own tmp directory, on the + // worker's disk and removed with the action, bound over /tmp (the + // bind carries the work volume's mount options, noexec and the + // like, not the host /tmp's; the config doc says so) so + // concurrent actions cannot collide on predictable paths there and + // nothing leaks between actions through it. A directory on disk + // rather than a tmpfs, because tmpfs pages are memory charged to + // the action's cgroup: a tool filling a memory-backed /tmp would + // take the worker down with it, where filling a directory under + // the scratch volume is bounded by that volume and the disk guard. + // This has to happen before the root action directory is masked: if + // that directory lives under /tmp the bind hides it, so its path is + // recreated inside the bound directory and the action directory + // itself is restored below from the saved file descriptor. When the + // root action directory lives elsewhere every component already + // exists and the mkdir is a no-op. + bind_mount(tmp_directory, c"/tmp")?; + mkdir_p_signal_safe(root_action_directory)?; + } + + // Mask the root action directory with a tmpfs to ensure sibling + // directories aren't visible. With the root at or under a bound /tmp + // the bind already hides every sibling, and a mask there would either + // land on top of the bind (root exactly /tmp: the action's "private + // /tmp" would be the 1 MiB mask) or leave a mount point sitting in the + // action's otherwise empty /tmp, so it is skipped. + let root_under_bound_tmp = tmp_directory.is_some() && is_under_tmp(root_action_directory); + if !root_under_bound_tmp { + let options = tmp_directory.is_some().then_some(ROOT_MASK_TMPFS_OPTIONS); + mount_tmpfs(root_action_directory, options)?; } // Recreate the specific operation's directory inside the empty tmpfs. @@ -396,10 +571,17 @@ fn perform_remount( /// SIGKILL to the actual process in the new user, PID, UTS and IPC /// namespaces. Pass this function to `CommandBuilder::pre_exec`. /// +/// When `mount` is set the process also gets its own mount namespace in +/// which the root action directory is masked so sibling actions are not +/// visible. When `tmp_directory` is given as well, that directory (the +/// action's own, already created) is bound over `/tmp`, so the action has +/// a private `/tmp` that is removed with it. +/// /// This function is async-signal-safe and has no external locks or /// memory allocations. pub fn configure_namespace( mount: bool, + tmp_directory: Option<&core::ffi::CStr>, root_action_directory: &core::ffi::CStr, action_directory: &core::ffi::CStr, ) -> std::io::Result<()> { @@ -444,7 +626,7 @@ pub fn configure_namespace( // Configure the mount namespace if enabled. if mount { - perform_remount(root_action_directory, action_directory).unwrap(); + perform_remount(tmp_directory, root_action_directory, action_directory)?; } // Set hostname to "nativelink" to ensure reproducibility. @@ -529,8 +711,16 @@ pub fn configure_namespace( let sig = unsafe { libc::sigwaitinfo(sigset.as_ptr(), siginfo.as_mut_ptr()) }; if sig == libc::SIGTERM { + // Pass the SIGTERM on, so the action gets the grace the + // worker's kill_grace_ms promises to write its own + // cleanup. The stub used to turn it into SIGKILL, which + // ended the action at once under namespaces while the + // same action got its grace without them. The hard kill + // is a SIGKILL to the stub itself: as this PID + // namespace's init, its death takes every process in + // the namespace with it. // SAFETY: pid is valid and we are sending a signal. - unsafe { libc::kill(pid, libc::SIGKILL) }; + unsafe { libc::kill(pid, libc::SIGTERM) }; } } } diff --git a/nativelink-worker/src/persistent_worker/live_worker.rs b/nativelink-worker/src/persistent_worker/live_worker.rs index b811fd6926f..e54f0d5a7b1 100644 --- a/nativelink-worker/src/persistent_worker/live_worker.rs +++ b/nativelink-worker/src/persistent_worker/live_worker.rs @@ -54,6 +54,8 @@ pub struct LiveWorker { /// Spawned child process. Kept on the struct so its Drop kills the process /// when the worker is dropped without `shutdown` being called. child: Child, + /// Registered with the reaper for as long as the child is ours to wait on. + _owned: crate::reaper::OwnedChild, stdin: ChildStdin, stdout: BufReader, /// Wire format we negotiated at spawn (immutable for this worker's lifetime). @@ -80,11 +82,17 @@ impl LiveWorker { startup_args: &[String], wire_format: WireFormat, working_dir: &Path, + env: &[(String, String)], + namespaced: bool, ) -> Result { let mut cmd = Command::new(executable); cmd.args(startup_args) .arg("--persistent_worker") .current_dir(working_dir) + // The action's environment and nothing else, as for a one-shot + // action; the environment is part of the worker's key. + .env_clear() + .envs(env.iter().map(|(name, value)| (name, value))) .stdin(Stdio::piped()) .stdout(Stdio::piped()) // WorkResponse.output carries per-action diagnostics. The child @@ -92,6 +100,30 @@ impl LiveWorker { // safely to a single request. .stderr(Stdio::null()) .kill_on_drop(true); + // Its own process group, so the resource sampler can attribute the + // process and its children to the request being served. + #[cfg(unix)] + cmd.process_group(0); + #[cfg(target_os = "linux")] + if namespaced { + use std::os::unix::ffi::OsStrExt; + let working_dir_c = std::ffi::CString::new(working_dir.as_os_str().as_bytes()) + .err_tip(|| "Persistent worker working directory is not a valid C string")?; + let action_dir_c = working_dir_c.clone(); + // SAFETY: configure_namespace is async-signal-safe and meant for pre_exec. + unsafe { + cmd.pre_exec(move || { + crate::namespace_utils::configure_namespace( + false, + None, + &working_dir_c, + &action_dir_c, + ) + }); + } + } + #[cfg(not(target_os = "linux"))] + let _ = namespaced; debug!( ?executable, @@ -106,6 +138,9 @@ impl LiveWorker { executable.display() ) })?; + // Ours until shutdown waits on it, however long it idles; the + // reaper leaves it alone. + let owned = crate::reaper::OwnedChild::new(child.id()); let stdin = child .stdin .take() @@ -117,6 +152,7 @@ impl LiveWorker { Ok(Self { child, + _owned: owned, stdin, stdout: BufReader::new(stdout), wire_format, @@ -130,6 +166,11 @@ impl LiveWorker { self.wire_format } + /// The worker process's id, which is also its process group. + pub fn pid(&self) -> Option { + self.child.id() + } + pub const fn request_count(&self) -> u64 { self.request_count } @@ -342,6 +383,21 @@ mod tests { } } + /// The worker process gets only the environment it is given. The shell + /// scripts need PATH to find `sleep` wherever the tests run; PowerShell + /// needs the rest of what Windows sets, so there the whole environment + /// goes through, as `action_environment` does with its defaults. + fn path_env() -> Vec<(String, String)> { + if cfg!(windows) { + std::env::vars().collect() + } else { + vec![( + "PATH".to_string(), + std::env::var("PATH").unwrap_or_default(), + )] + } + } + #[cfg(unix)] fn echo_script(working_dir: &Path, unix_body: &str, _windows_body: &str) -> TestWorkerProgram { let path = working_dir.join("worker.sh"); @@ -386,6 +442,8 @@ mod tests { script.startup_args(), WireFormat::Json, dir.path(), + &path_env(), + false, ) .unwrap(); let start = Instant::now(); @@ -409,6 +467,8 @@ mod tests { script.startup_args(), WireFormat::Json, dir.path(), + &path_env(), + false, ) .unwrap(); @@ -438,6 +498,8 @@ mod tests { script.startup_args(), WireFormat::Json, dir.path(), + &path_env(), + false, ) .unwrap(); @@ -466,6 +528,8 @@ mod tests { script.startup_args(), WireFormat::Json, dir.path(), + &path_env(), + false, ) .unwrap(); let req = WorkRequest { diff --git a/nativelink-worker/src/persistent_worker/pool.rs b/nativelink-worker/src/persistent_worker/pool.rs index 74840ef98ce..444a7b6a542 100644 --- a/nativelink-worker/src/persistent_worker/pool.rs +++ b/nativelink-worker/src/persistent_worker/pool.rs @@ -19,15 +19,17 @@ //! executable + startup-flag prefix + wire format are identical share a worker //! process. -use core::mem; use core::time::Duration; +use core::{fmt, mem}; use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::sync::Arc; use nativelink_error::{Code, Error, ResultExt, make_err}; use nativelink_util::background_spawn; +use nativelink_util::task::JoinHandleDropGuard; use parking_lot::Mutex; +use tokio::sync::Notify; use tracing::{debug, info, warn}; use super::live_worker::LiveWorker; @@ -45,9 +47,12 @@ pub const DEFAULT_IDLE_TIMEOUT: Duration = Duration::from_mins(5); /// dropped on `release` (helps recycle accumulated JVM state). pub const DEFAULT_MAX_REQUESTS_PER_WORKER: u64 = 200; +/// Default wait for a worker to come back when a key is at its cap. +pub const DEFAULT_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(30); + /// Identity by which two persistent-worker actions are considered compatible. /// Same `WorkerKey` => same worker can serve both. -#[derive(Clone, Debug, PartialEq, Eq, Hash)] +#[derive(Clone, PartialEq, Eq, Hash)] pub struct WorkerKey { /// Resolved tool executable. Mirrors `Action.arguments[0]`. pub executable: PathBuf, @@ -57,6 +62,27 @@ pub struct WorkerKey { pub startup_args: Vec, /// Wire format from the action's `requires-worker-protocol`. pub wire_format: WireFormat, + /// The action's environment, sorted by name. A worker process is + /// started with exactly this environment, the way a one-shot action is, + /// so two actions that differ in it do not share a process. + pub env: Vec<(String, String)>, +} + +/// The key is logged on every spawn and completion and quoted in errors the +/// client sees, so the environment shows as its variable names only: an +/// action's environment can carry a token. +impl fmt::Debug for WorkerKey { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("WorkerKey") + .field("executable", &self.executable) + .field("startup_args", &self.startup_args) + .field("wire_format", &self.wire_format) + .field( + "env", + &self.env.iter().map(|(name, _)| name).collect::>(), + ) + .finish() + } } impl WorkerKey { @@ -83,8 +109,17 @@ impl WorkerKey { executable: PathBuf::from(exe), startup_args, wire_format, + env: Vec::new(), }) } + + /// The environment the worker process is started with, part of the key. + #[must_use] + pub fn with_env(mut self, mut env: Vec<(String, String)>) -> Self { + env.sort(); + self.env = env; + self + } } /// Configuration knobs for a `PersistentWorkerPool`. Defaults are sensible for @@ -96,6 +131,16 @@ pub struct PoolConfig { pub max_requests_per_worker: u64, /// SIGKILL grace when shutting a worker down. pub shutdown_grace: Duration, + /// How long `acquire` waits for a worker to come back when the key is + /// at `max_workers_per_key` with none idle, before the action runs + /// one-shot instead. Bazel itself queues on `worker_max_instances`; a + /// one-shot run of a worker tool pays the tool's startup again. + pub acquire_timeout: Duration, + /// Start each worker process in its own user, PID, UTS and IPC + /// namespaces, as one-shot actions are when the worker runs with + /// `use_namespaces`. Never a mount namespace: the process outlives any + /// one action, so a private `/tmp` per action cannot apply to it. + pub namespaced: bool, } impl Default for PoolConfig { @@ -105,6 +150,8 @@ impl Default for PoolConfig { idle_timeout: DEFAULT_IDLE_TIMEOUT, max_requests_per_worker: DEFAULT_MAX_REQUESTS_PER_WORKER, shutdown_grace: Duration::from_secs(5), + acquire_timeout: DEFAULT_ACQUIRE_TIMEOUT, + namespaced: false, } } } @@ -230,15 +277,21 @@ impl KeyState { struct PoolInner { config: PoolConfig, state: Mutex>, + /// Woken whenever a worker is returned or a slot is freed, for an + /// `acquire` waiting on a key at its cap. + returned: Notify, } impl PoolInner { /// Helper used by Drop on a forgotten Lease. fn decrement_count(&self, key: &WorkerKey) { - let mut state = self.state.lock(); - if let Some(s) = state.get_mut(key) { - s.total_count = s.total_count.saturating_sub(1); + { + let mut state = self.state.lock(); + if let Some(s) = state.get_mut(key) { + s.total_count = s.total_count.saturating_sub(1); + } } + self.returned.notify_waiters(); } async fn return_worker(self: Arc, key: WorkerKey, worker: LiveWorker, healthy: bool) { @@ -255,9 +308,12 @@ impl PoolInner { count_slot.decrement_now(); return; } - let mut state = self.state.lock(); - let entry = state.entry(key).or_insert_with(KeyState::new); - entry.idle.push(worker); + { + let mut state = self.state.lock(); + let entry = state.entry(key).or_insert_with(KeyState::new); + entry.idle.push(worker); + } + self.returned.notify_waiters(); } } @@ -274,11 +330,17 @@ impl Default for PersistentWorkerPool { } impl PersistentWorkerPool { + /// How long `acquire` waits at the cap before giving up. + pub fn acquire_timeout(&self) -> Duration { + self.inner.config.acquire_timeout + } + pub fn new(config: PoolConfig) -> Self { Self { inner: Arc::new(PoolInner { config, state: Mutex::new(HashMap::new()), + returned: Notify::new(), }), } } @@ -298,34 +360,49 @@ impl PersistentWorkerPool { executable_path: &Path, working_dir: &Path, ) -> Result { - // First, try to take an idle worker without blocking on spawn. - { - let mut state = self.inner.state.lock(); - let entry = state.entry(key.clone()).or_insert_with(KeyState::new); - while let Some(mut worker) = entry.idle.pop() { - if worker.is_dead() { - // Process died while idle (e.g. JVM OOM). Drop it. - entry.total_count = entry.total_count.saturating_sub(1); - continue; + let deadline = tokio::time::Instant::now() + self.inner.config.acquire_timeout; + loop { + // Take an idle worker, or reserve a slot to spawn one, without + // blocking on the spawn itself. + let notified = { + let mut state = self.inner.state.lock(); + let entry = state.entry(key.clone()).or_insert_with(KeyState::new); + while let Some(mut worker) = entry.idle.pop() { + if worker.is_dead() { + // Process died while idle (e.g. JVM OOM). Drop it. + entry.total_count = entry.total_count.saturating_sub(1); + continue; + } + return Ok(Lease { + inner: Some(LeaseInner { + worker, + pool: self.inner.clone(), + key, + }), + }); } - return Ok(Lease { - inner: Some(LeaseInner { - worker, - pool: self.inner.clone(), - key, - }), - }); - } - if entry.total_count >= self.inner.config.max_workers_per_key { + if entry.total_count < self.inner.config.max_workers_per_key { + // Reserve a slot before releasing the lock so a concurrent + // acquire can't over-commit. + entry.total_count += 1; + break; + } + // At the cap with nothing idle: wait for a worker to come + // back. The waiter is registered under the lock so a return + // between the check and the wait is not missed. + let notified = self.inner.returned.notified(); + let mut notified = Box::pin(notified); + notified.as_mut().enable(); + notified + }; + if tokio::time::timeout_at(deadline, notified).await.is_err() { return Err(make_err!( Code::ResourceExhausted, - "Persistent worker pool for {key:?} is at capacity ({} workers)", - self.inner.config.max_workers_per_key + "Persistent worker pool for {key:?} is at capacity ({} workers) and none came back within {:?}", + self.inner.config.max_workers_per_key, + self.inner.config.acquire_timeout )); } - // Reserve a slot before releasing the lock so a concurrent - // acquire can't over-commit. - entry.total_count += 1; } // Lock released; spawn outside the critical section. @@ -335,6 +412,8 @@ impl PersistentWorkerPool { &key.startup_args, key.wire_format, working_dir, + &key.env, + self.inner.config.namespaced, ); let worker = match spawn_result { Ok(w) => w, @@ -353,6 +432,32 @@ impl PersistentWorkerPool { }) } + /// Runs `sweep_idle` on a timer until the pool is dropped. Without it an + /// idle worker process, a JVM as a rule, lived as long as the worker did. + pub fn spawn_sweeper(&self) -> JoinHandleDropGuard<()> { + let inner = Arc::downgrade(&self.inner); + let interval = (self.inner.config.idle_timeout / 4).max(Duration::from_secs(1)); + JoinHandleDropGuard::new(background_spawn!("persistent_worker_sweeper", async move { + loop { + tokio::time::sleep(interval).await; + let Some(inner) = inner.upgrade() else { + return; + }; + Self { inner }.sweep_idle().await; + } + })) + } + + /// How many worker processes the pool holds, idle or leased. + pub fn live_workers(&self) -> usize { + self.inner + .state + .lock() + .values() + .map(|entry| entry.total_count) + .sum() + } + /// Drop all idle workers whose `last_used` is older than `idle_timeout`. /// Intended to be called periodically from a background task. pub async fn sweep_idle(&self) { @@ -397,12 +502,23 @@ mod tests { path } + /// The worker process gets only the environment in its key, so the + /// scripts need PATH to find `sleep` wherever the tests run. + #[cfg(unix)] + fn path_env() -> Vec<(String, String)> { + vec![( + "PATH".to_string(), + std::env::var("PATH").unwrap_or_default(), + )] + } + #[cfg(unix)] fn shell_worker_key(script: &Path) -> WorkerKey { WorkerKey { executable: PathBuf::from("/bin/sh"), startup_args: vec![script.display().to_string()], wire_format: WireFormat::Json, + env: path_env(), } } @@ -484,6 +600,52 @@ mod tests { assert_ne!(a, b); } + #[test] + fn worker_key_distinguishes_environments_and_ignores_their_order() { + let base = + || WorkerKey::from_argv(&["javac", "@a"].map(String::from), WireFormat::Proto).unwrap(); + let a = base().with_env(vec![ + ("PATH".into(), "/bin".into()), + ("LANG".into(), "C".into()), + ]); + let b = base().with_env(vec![ + ("LANG".into(), "C".into()), + ("PATH".into(), "/bin".into()), + ]); + let c = base().with_env(vec![("PATH".into(), "/usr/bin".into())]); + assert_eq!(a, b); + assert_ne!(a, c); + } + + #[nativelink_test] + #[cfg(unix)] + async fn sweeper_retires_a_worker_idle_past_the_timeout() { + let dir = tempfile::tempdir().unwrap(); + let script = shell_script(dir.path(), "exec sleep 60\n"); + let key = shell_worker_key(&script); + let pool = PersistentWorkerPool::new(PoolConfig { + idle_timeout: Duration::from_millis(200), + shutdown_grace: Duration::from_millis(100), + ..PoolConfig::default() + }); + let _sweeper = pool.spawn_sweeper(); + + let lease = pool + .acquire(key, Path::new("/bin/sh"), dir.path()) + .await + .unwrap(); + lease.release(true).await; + assert_eq!(pool.live_workers(), 1); + + // The sweeper runs every second at least; by then the worker has + // been idle far longer than its timeout. + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + while pool.live_workers() > 0 && tokio::time::Instant::now() < deadline { + tokio::time::sleep(Duration::from_millis(100)).await; + } + assert_eq!(pool.live_workers(), 0, "the idle worker was not retired"); + } + #[nativelink_test] #[cfg(unix)] async fn dropping_unhealthy_release_frees_pool_slot() { @@ -516,6 +678,45 @@ mod tests { next_lease.release(false).await; } + #[nativelink_test] + #[cfg(unix)] + async fn acquire_at_the_cap_waits_for_a_returned_worker() { + let dir = tempfile::tempdir().unwrap(); + let script = shell_script(dir.path(), "exec sleep 60\n"); + let key = shell_worker_key(&script); + let pool = PersistentWorkerPool::new(PoolConfig { + max_workers_per_key: 1, + shutdown_grace: Duration::from_millis(100), + acquire_timeout: Duration::from_secs(5), + ..PoolConfig::default() + }); + + let lease = pool + .acquire(key.clone(), Path::new("/bin/sh"), dir.path()) + .await + .unwrap(); + let waiter = { + let pool = pool.clone(); + let key = key.clone(); + let dir = dir.path().to_path_buf(); + background_spawn!("persistent_worker_acquire_waiter", async move { + pool.acquire(key, Path::new("/bin/sh"), &dir).await + }) + }; + tokio::time::sleep(Duration::from_millis(100)).await; + assert!( + !waiter.is_finished(), + "the second acquire should be waiting at the cap" + ); + lease.release(true).await; + let second = tokio::time::timeout(Duration::from_secs(2), waiter) + .await + .expect("the waiter got the returned worker") + .unwrap() + .unwrap(); + second.release(false).await; + } + #[nativelink_test] #[cfg(unix)] async fn concurrent_acquire_respects_per_key_cap() { @@ -525,6 +726,7 @@ mod tests { let pool = PersistentWorkerPool::new(PoolConfig { max_workers_per_key: 1, shutdown_grace: Duration::from_millis(100), + acquire_timeout: Duration::from_millis(200), ..PoolConfig::default() }); diff --git a/nativelink-worker/src/reaper.rs b/nativelink-worker/src/reaper.rs new file mode 100644 index 00000000000..f09e6f01df4 --- /dev/null +++ b/nativelink-worker/src/reaper.rs @@ -0,0 +1,172 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! What an action leaves behind is the worker's to reap. +//! +//! The worker waits on the one process it spawned per action. Anything that +//! process forked and did not wait for (a shell's background job, a helper +//! that daemonized) is reparented when it exits: to the worker when the +//! worker is the container's PID 1, which it usually is, and otherwise to +//! whatever init the container has. The worker never waited on those, so +//! each became a zombie the moment it exited and stayed one for the life of +//! the pod; a worker seen with hundreds of them had run one build whose +//! actions forked and did not wait. This module makes the worker their +//! parent wherever it runs and reaps them on a timer. +//! +//! Inside `use_namespaces` the stub is the init of the action's PID +//! namespace and reaps there; this covers the actions that run without it. + +use std::collections::{BTreeSet, HashSet}; +use std::sync::Mutex; + +/// How often the worker reaps the zombies its actions left behind. +pub const REAP_INTERVAL: core::time::Duration = core::time::Duration::from_secs(5); + +/// The pids of every child this process spawned and still holds a handle +/// for. The reaper never touches these: their exit status belongs to the +/// code that waits on them, however long it takes to get round to it (a +/// persistent worker is only `try_wait`ed on demand, an action's wait can +/// be starved). Ownership is decided here, not by timing. +static OWNED: Mutex>> = Mutex::new(None); + +/// A child this process owns, for as long as the guard lives. +#[derive(Debug)] +pub struct OwnedChild(Option); + +impl OwnedChild { + /// Registers `pid` as ours; `None` (a child that already exited at + /// spawn) registers nothing. Drop the guard once the child has been + /// waited on, or with the handle that will wait on it. + pub fn new(pid: Option) -> Self { + if let Some(pid) = pid { + OWNED + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get_or_insert_with(HashSet::new) + .insert(pid); + } + Self(pid) + } +} + +impl Drop for OwnedChild { + fn drop(&mut self) { + if let Some(pid) = self.0 + && let Some(owned) = OWNED + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .as_mut() + { + owned.remove(&pid); + } + } +} + +#[cfg(target_os = "linux")] +fn is_owned(pid: u32) -> bool { + OWNED + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .as_ref() + .is_some_and(|owned| owned.contains(&pid)) +} + +/// Make this process the parent of every descendant an action leaves +/// behind, even when it is not PID 1, so the zombies land here where they +/// can be reaped rather than with an init that never will. +#[cfg(target_os = "linux")] +pub fn become_subreaper() -> Result<(), std::io::Error> { + // SAFETY: prctl with PR_SET_CHILD_SUBREAPER takes integers only and + // changes a flag on this process. + if unsafe { libc::prctl(libc::PR_SET_CHILD_SUBREAPER, 1, 0, 0, 0) } == 0 { + Ok(()) + } else { + Err(std::io::Error::last_os_error()) + } +} + +#[cfg(not(target_os = "linux"))] +pub const fn become_subreaper() -> Result<(), std::io::Error> { + Ok(()) +} + +/// The parent pid out of a `/proc//stat` line, when the process is a +/// zombie; `None` for a live process or unreadable text. `comm` can hold +/// spaces and parentheses, so the fields are read after the final `)`: +/// state first, then the parent pid. +pub fn parse_zombie_ppid(stat: &str) -> Option { + let after_comm = stat.rsplit_once(')')?.1; + let mut fields = after_comm.split_whitespace(); + let state = fields.next()?; + let ppid = fields.next()?.parse().ok()?; + (state == "Z").then_some(ppid) +} + +/// The zombies whose parent is this process, by pid. One short read per +/// process; call it off the runtime, since `/proc` can be large. +#[cfg(target_os = "linux")] +pub fn zombie_children() -> BTreeSet { + // SAFETY: getpid has no memory safety considerations. + let me = u32::try_from(unsafe { libc::getpid() }).unwrap_or(0); + let Ok(entries) = std::fs::read_dir("/proc") else { + return BTreeSet::new(); + }; + entries + .flatten() + .filter_map(|entry| entry.file_name().to_string_lossy().parse::().ok()) + .filter(|pid| { + std::fs::read_to_string(format!("/proc/{pid}/stat")) + .ok() + .and_then(|stat| parse_zombie_ppid(&stat)) + == Some(me) + }) + .collect() +} + +/// Reaps the zombies that are not ours to wait on and were already zombies +/// at the previous sweep; returns how many, and this sweep's zombies for +/// the next call. A child this process spawned is never touched, whatever +/// its state: its exit status belongs to whoever holds its handle. The +/// second sweep is a guard on top of that, for a pid we never registered +/// (a fork of ours outside these paths); a pid reused within an interval +/// starts over, since a reaped pid is not carried forward. Blocking; run it +/// on the blocking pool. +#[cfg(target_os = "linux")] +pub fn reap_orphaned_zombies(seen_last: &BTreeSet) -> (usize, BTreeSet) { + let mut now = zombie_children(); + let mut reaped = 0; + let candidates: Vec = now + .iter() + .copied() + .filter(|pid| seen_last.contains(pid) && !is_owned(*pid)) + .collect(); + for pid in candidates { + let Ok(pid_t) = libc::pid_t::try_from(pid) else { + continue; + }; + let mut status = 0; + // SAFETY: waitpid takes a pid and a pointer to an int on the stack; + // a pid that is not our child is reported as ECHILD, not acted on. + if unsafe { libc::waitpid(pid_t, &raw mut status, libc::WNOHANG) } == pid_t { + reaped += 1; + now.remove(&pid); + } + } + (reaped, now) +} + +#[cfg(not(target_os = "linux"))] +pub const fn reap_orphaned_zombies(_seen_last: &BTreeSet) -> (usize, BTreeSet) { + (0, BTreeSet::new()) +} diff --git a/nativelink-worker/src/running_actions_manager.rs b/nativelink-worker/src/running_actions_manager.rs index 57ce6e7e1e9..9c76c6e17d6 100644 --- a/nativelink-worker/src/running_actions_manager.rs +++ b/nativelink-worker/src/running_actions_manager.rs @@ -20,13 +20,15 @@ use core::sync::atomic::{AtomicBool, Ordering}; use core::time::Duration; use std::borrow::Cow; use std::collections::vec_deque::VecDeque; -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::env; use std::ffi::{OsStr, OsString}; #[cfg(target_family = "unix")] use std::fs::Permissions; #[cfg(target_family = "unix")] use std::os::unix::fs::{MetadataExt, PermissionsExt}; +#[cfg(target_family = "unix")] +use std::os::unix::process::ExitStatusExt; use std::path::{Path, PathBuf}; use std::process::Stdio; use std::sync::{Arc, Weak}; @@ -36,11 +38,12 @@ use bytes::{Bytes, BytesMut}; use filetime::{FileTime, set_file_mtime}; use formatx::Template; use futures::future::{ - BoxFuture, Future, FutureExt, TryFutureExt, join_all, try_join, try_join_all, + BoxFuture, Fuse, Future, FutureExt, TryFutureExt, join_all, try_join, try_join_all, }; use futures::stream::{FuturesUnordered, StreamExt, TryStreamExt}; use nativelink_config::cas_server::{ - Buck2FileCaptureConfig, EnvironmentSource, UploadActionResultConfig, UploadCacheResultsStrategy, + Buck2FileCaptureConfig, DiskEnforcement, EnvironmentSource, MemoryEnforcement, + ResourceEnforcementConfig, UploadActionResultConfig, UploadCacheResultsStrategy, }; use nativelink_error::{Code, Error, ResultExt, make_err, make_input_err}; use nativelink_metric::MetricsComponent; @@ -50,7 +53,7 @@ use nativelink_proto::build::bazel::remote::execution::v2::{ Tree as ProtoTree, UpdateActionResultRequest, }; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{ - ActionResourceUsage, HistoricalExecuteResponse, StartExecute, + ActionResourceUsage, HistoricalExecuteResponse, Reservation, ResourceOutcome, StartExecute, }; use nativelink_store::ac_utils::{ ESTIMATED_DIGEST_SIZE, compute_buf_digest, get_and_decode_digest, serialize_and_upload_message, @@ -67,6 +70,7 @@ use nativelink_util::common::{DigestInfo, fs}; use nativelink_util::digest_hasher::{DigestHasher, DigestHasherFunc}; use nativelink_util::metrics_utils::{AsyncCounterWrapper, CounterWithTime}; use nativelink_util::store_trait::{Store, StoreLike, UploadSizeInfo}; +use nativelink_util::task::JoinHandleDropGuard; use nativelink_util::{background_spawn, spawn, spawn_blocking}; use parking_lot::Mutex; use prost::Message; @@ -75,7 +79,7 @@ use scopeguard::{ScopeGuard, guard}; use serde::Deserialize; use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt}; use tokio::process; -use tokio::sync::{Notify, oneshot, watch}; +use tokio::sync::{Notify, Semaphore, oneshot, watch}; use tokio::time::Instant; use tokio_stream::wrappers::ReadDirStream; use tonic::Request; @@ -84,7 +88,8 @@ use uuid::Uuid; use crate::buck2_file_capture::Buck2FileCapture; use crate::persistent_worker::{ - Input as PersistentWorkerInput, PersistentWorkerPool, WireFormat, WorkRequest, WorkerKey, + Input as PersistentWorkerInput, PersistentWorkerPool, PoolConfig, WireFormat, WorkRequest, + WorkerKey, }; /// For simplicity we use a fixed exit code for cases when our program is terminated @@ -103,12 +108,150 @@ const DEFAULT_HISTORICAL_RESULTS_STRATEGY: UploadCacheResultsStrategy = #[cfg(target_os = "linux")] const RESOURCE_USAGE_SAMPLE_INTERVAL: Duration = Duration::from_millis(250); +/// How many output files are open for upload at once, across the whole +/// output tree: one semaphore, taken around each file from open to upload, +/// so nested directories cannot multiply the bound. An unbounded set has hit +/// the descriptor limit on actions with tens of thousands of outputs. The +/// same number bounds each stream's fan-out. +const UPLOAD_CONCURRENCY: usize = 64; + /// What the sampler observed over an action's lifetime. -#[cfg(target_os = "linux")] #[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] struct SampledResourceUsage { peak_memory_kb: u64, cpu_time_ms: u64, + /// The last reading before the group was gone, which is what a kernel + /// kill is judged against. + last_memory_kb: u64, + /// The most the action's own files under its directory came to, when + /// disk was measured; 0 otherwise. + peak_disk_kb: u64, +} + +/// Why the worker ended an action before the action ended itself. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum KillReason { + Timeout, + Memory, + Disk, + External, +} + +const SIGKILL_NUMBER: i32 = 9; + +/// Appended to a `NotFound` raised while fetching something the client +/// referenced by digest (the action, its command, its inputs), so the +/// worker loop can tell a missing input from any other missing thing +/// without reading a store's wording. +pub const MISSING_INPUT_ERROR_TIP: &str = "action input missing from CAS"; + +fn tag_missing_input(err: Error) -> Error { + if err.code == Code::NotFound { + err.append(MISSING_INPUT_ERROR_TIP) + } else { + err + } +} + +/// How long a killed action's output pipes, or those of one that left +/// stragglers behind, get to drain. They close when every process holding +/// them is gone, which the SIGKILL to the group sees to; a process that +/// left the group could still hold them, so that wait is bounded, and what +/// was read by then is kept. An action that exited cleanly with nothing +/// left behind is read to the end. +const OUTPUT_DRAIN_TIMEOUT: Duration = Duration::from_secs(5); + +/// Classifies how an action ended. A kill the worker sent names itself; a +/// SIGKILL nobody here sent, with the last memory sample within 10% of the +/// limit, is the kernel's OOM killer; any other signal is external. +pub fn classify_outcome( + kill_reason: Option, + signal: Option, + last_memory_kb: u64, + memory_limit_kb: Option, +) -> (ResourceOutcome, bool) { + match (kill_reason, signal) { + (Some(KillReason::Memory), _) => (ResourceOutcome::KilledMemory, true), + (Some(KillReason::Disk), _) => (ResourceOutcome::KilledDisk, true), + (Some(KillReason::Timeout), _) => (ResourceOutcome::KilledTimeout, false), + (None, None) => (ResourceOutcome::Completed, false), + (None, Some(SIGKILL_NUMBER)) => { + let near_limit = memory_limit_kb.is_some_and(|limit_kb| { + limit_kb > 0 && last_memory_kb.saturating_mul(10) >= limit_kb.saturating_mul(9) + }); + if near_limit { + (ResourceOutcome::KilledMemory, false) + } else { + (ResourceOutcome::KilledExternal, false) + } + } + (Some(KillReason::External), _) | (None, Some(_)) => { + (ResourceOutcome::KilledExternal, false) + } + } +} + +#[cfg(target_os = "linux")] +type SpawnedChild = crate::namespace_utils::MaybeNamespacedChild; +#[cfg(not(target_os = "linux"))] +type SpawnedChild = process::Child; + +/// SIGKILL through the child wrapper, which takes the whole process group +/// with it, logged if it fails. +async fn hard_kill(child: &mut SpawnedChild, why: &str) { + if let Err(err) = child.kill().await { + error!(?err, why, "Could not kill process in RunningActionsManager"); + } +} + +/// Starts stopping an action. With a grace period the process group gets +/// SIGTERM and the returned future fires when SIGKILL is due; without one +/// SIGKILL goes out now and the future never fires. +async fn stop_action( + child: &mut SpawnedChild, + grace: Duration, + why: &str, +) -> Fuse> { + debug!(grace_ms = grace.as_millis(), why, "Stopping action"); + #[cfg(target_os = "linux")] + if !grace.is_zero() + && let Some(pgid) = child.id() + { + signal_process_group(pgid, libc::SIGTERM); + return tokio::time::sleep(grace).boxed().fuse(); + } + hard_kill(child, why).await; + Fuse::terminated() +} + +/// The reservation the action was admitted with, read back from its +/// platform properties so the measurement travels with its yardstick. +fn reservation_of( + properties: &HashMap, + memory_property: Option<&str>, + disk_property: Option<&str>, +) -> Option { + let read = |name: &str| { + properties + .get(name) + .and_then(|v| v.trim().parse::().ok()) + }; + let reservation = Reservation { + cpu_count: read("cpu_count").unwrap_or(0), + memory_kb: read(memory_property.unwrap_or("memory_kb")).unwrap_or(0), + disk_kb: read(disk_property.unwrap_or("disk_kb")).unwrap_or(0), + }; + (reservation.cpu_count > 0 || reservation.memory_kb > 0 || reservation.disk_kb > 0) + .then_some(reservation) +} + +#[cfg(target_os = "linux")] +fn pod_memory_limit_kb() -> Option { + crate::capacity::observe_cgroup().map(|observed| observed.memory_kb) +} +#[cfg(not(target_os = "linux"))] +const fn pod_memory_limit_kb() -> Option { + None } #[cfg(target_os = "linux")] @@ -117,16 +260,148 @@ struct ActionResourceUsageSampler { handle: tokio::task::JoinHandle, } +/// What the sampler found over a ceiling, with the figure it saw. Only the +/// Linux sampler sends one; the arm that receives it is built everywhere. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[cfg_attr(not(target_os = "linux"), allow(dead_code))] +enum OverLimit { + Memory(u64), + Disk(u64), +} + +/// The ceilings the sampler enforces. Memory: two consecutive samples over +/// `memory_limit_kb`, or one at twice it. Disk: one sample of the action's +/// own files over `disk_limit_kb`. The first breach is sent on +/// `over_limit_tx` and the caller kills the action; nothing is sent when no +/// ceiling is set, and the caller keeps the sender alive so the receiver +/// stays pending. +#[cfg(target_os = "linux")] +struct Ceilings { + memory_limit_kb: Option, + disk_limit_kb: Option, + over_limit_tx: Option>, +} + +/// One disk sample every this many memory samples: a walk of the action +/// directory is a stat per file, so it runs every ten seconds, not four +/// times a second. +#[cfg(target_os = "linux")] +const DISK_SAMPLE_EVERY: u32 = 40; + #[cfg(target_os = "linux")] -fn start_action_resource_usage_sampler(pgid: u32) -> ActionResourceUsageSampler { +fn start_action_resource_usage_sampler( + pgid: u32, + ceilings: Ceilings, + disk_directory: Option<(PathBuf, SystemTime)>, +) -> ActionResourceUsageSampler { let (stop_tx, stop_rx) = watch::channel(false); let handle = background_spawn!( "action_resource_usage_sampler", - sample_action_resource_usage(pgid, stop_rx) + sample_action_resource_usage(pgid, stop_rx, ceilings, disk_directory) ); ActionResourceUsageSampler { stop_tx, handle } } +/// KiB on disk of the files under `directory` that the action wrote: regular +/// files modified at or after `since`, the directory's stamp from just +/// before its command started (see [`directory_stamp`]). Its +/// inputs were materialized before that, as hard links into the CAS store +/// or private copies, and keep their earlier times, so they are told apart +/// by time rather than by link count: a hard link drops to one link the +/// moment the store evicts the blob, and would then pass for the action's +/// own. Symlinks and the directories' own blocks are not counted. Blocking; +/// run it on the blocking pool. +#[cfg(target_family = "unix")] +pub fn directory_private_kb(directory: &Path, since: SystemTime) -> u64 { + let mut total_blocks: u64 = 0; + let mut pending = vec![directory.to_path_buf()]; + while let Some(dir) = pending.pop() { + let Ok(entries) = std::fs::read_dir(&dir) else { + continue; + }; + for entry in entries.flatten() { + let Ok(metadata) = entry.metadata() else { + continue; + }; + let file_type = metadata.file_type(); + if file_type.is_dir() { + pending.push(entry.path()); + } else if file_type.is_file() + && metadata.modified().is_ok_and(|modified| modified >= since) + { + total_blocks = total_blocks.saturating_add(metadata.blocks()); + } + } + } + // st_blocks are 512-byte units whatever the filesystem's block size. + total_blocks / 2 +} + +#[cfg(not(target_family = "unix"))] +pub fn directory_private_kb(_directory: &Path, _since: SystemTime) -> u64 { + 0 +} + +/// Touches `directory` and returns the time the filesystem recorded for it. +/// File mtimes come from the kernel's coarse clock, which on Linux can lag +/// the wall clock by a tick, so a file written right after +/// `SystemTime::now()` can look older than that moment. The walk compares +/// mtimes against this stamp, taken from the same clock, so a file the +/// action writes afterwards is never older than it. +#[cfg(target_family = "unix")] +pub fn directory_stamp(directory: &Path) -> Result { + use std::os::unix::io::AsRawFd; + let dir = std::fs::File::open(directory) + .err_tip(|| format!("Opening {} to stamp it", directory.display()))?; + // SAFETY: a zeroed timespec is valid; only tv_nsec is set. + let mut now: libc::timespec = unsafe { core::mem::zeroed() }; + now.tv_nsec = libc::UTIME_NOW; + let times = [now, now]; + // SAFETY: futimens takes an open fd and a pointer to two timespecs that + // outlive the call. + if unsafe { libc::futimens(dir.as_raw_fd(), times.as_ptr()) } != 0 { + return Err(std::io::Error::last_os_error()) + .err_tip(|| format!("Stamping {}", directory.display())); + } + dir.metadata() + .and_then(|metadata| metadata.modified()) + .err_tip(|| format!("Reading the stamp on {}", directory.display())) +} + +#[cfg(not(target_family = "unix"))] +pub fn directory_stamp(_directory: &Path) -> Result { + Ok(SystemTime::now()) +} + +/// Starts a disk sample: the action's own files, walked on the blocking +/// pool as a task of its own, so the memory samples carry on while a large +/// tree is being walked. `None` when there is no directory to measure. +#[cfg(target_os = "linux")] +fn start_disk_sample( + directory: Option<&(PathBuf, SystemTime)>, +) -> Option> { + let (directory, since) = directory?.clone(); + Some(spawn_blocking!("action_disk_sample", move || { + directory_private_kb(&directory, since) + })) +} + +/// A signal to every process in the action's group. The action is its own +/// group leader, so this reaches the children a shell forked. The SIGKILL +/// at the end of a kill goes through the child wrapper, which takes the +/// group with it (or tells the stub, namespaced); this is for the SIGTERM +/// that gives an action its grace, and for the stragglers left once the +/// action's own process has exited. Returns whether any process got it. +#[cfg(target_family = "unix")] +fn signal_process_group(pgid: u32, signal: i32) -> bool { + let Ok(pgid) = i32::try_from(pgid) else { + return false; + }; + // SAFETY: killpg only takes integers and has no memory safety + // considerations; a stale group id is reported as ESRCH, not acted on. + unsafe { libc::killpg(pgid, signal) == 0 } +} + #[cfg(target_os = "linux")] async fn finish_action_resource_usage_sampler( sampler: ActionResourceUsageSampler, @@ -139,13 +414,25 @@ async fn finish_action_resource_usage_sampler( async fn sample_action_resource_usage( pgid: u32, mut stop_rx: watch::Receiver, + mut ceilings: Ceilings, + disk_directory: Option<(PathBuf, SystemTime)>, ) -> SampledResourceUsage { let mut peak_memory_kb = 0; + let mut last_memory_kb = 0; + let mut peak_disk_kb = 0; + let mut samples_over_limit = 0u32; + let mut samples_since_disk = DISK_SAMPLE_EVERY; + // The walk in flight, if one is; polled, never awaited, from the loop. + let mut disk_sample: Option> = None; // CPU time is cumulative per process and a process that exits stops // appearing, so summing the live group at the end would lose everything // short-lived. Keep the last figure seen for each pid and total them at // the end instead. let mut cpu_ticks_by_pid = HashMap::new(); + // What the group had already spent when sampling began. A one-shot + // process starts at nothing; a pooled process carries every request it + // served before, which is not this action's. + let mut baseline_ticks: Option = None; let sample = |peak_memory_kb: &mut u64, cpu_ticks_by_pid: &mut HashMap| { let observed = sample_process_group(pgid, cpu_ticks_by_pid); @@ -156,13 +443,66 @@ async fn sample_action_resource_usage( }; loop { - if sample(&mut peak_memory_kb, &mut cpu_ticks_by_pid).is_none() - && !Path::new(&format!("/proc/{pgid}")).exists() - { + let observed = sample(&mut peak_memory_kb, &mut cpu_ticks_by_pid); + if baseline_ticks.is_none() { + baseline_ticks = Some(cpu_ticks_by_pid.values().sum()); + } + if let Some(memory_kb) = observed { + last_memory_kb = memory_kb; + } + if observed.is_none() && !Path::new(&format!("/proc/{pgid}")).exists() { // The group leader has been reaped and no member process remains, // so the action is finished. break; } + // Two consecutive samples over the ceiling, so one spike between + // reads is not a kill; one sample at twice the ceiling is, because + // an allocation that fast reaches the pod's limit before the next + // read (observed: twelve 20 GiB allocators filled a 52 GiB pod in + // under a second). + if let (Some(observed_kb), Some(limit_kb)) = (observed, ceilings.memory_limit_kb) { + if observed_kb > limit_kb { + samples_over_limit += 1; + } else { + samples_over_limit = 0; + } + let gross = observed_kb >= limit_kb.saturating_mul(2); + if (samples_over_limit >= 2 || gross) + && let Some(over_limit_tx) = ceilings.over_limit_tx.take() + && over_limit_tx.send(OverLimit::Memory(observed_kb)).is_err() + { + // The receiver is gone only when the action already ended. + debug!(observed_kb, "Memory ceiling breached, action already ended"); + } + } + + // A walk that finished since the last sample; the memory samples + // above never waited for it. + if let Some(handle) = disk_sample.as_mut() + && let core::task::Poll::Ready(finished) = futures::poll!(handle) + { + disk_sample = None; + if let Ok(disk_kb) = finished { + peak_disk_kb = peak_disk_kb.max(disk_kb); + if let Some(limit_kb) = ceilings.disk_limit_kb + && disk_kb > limit_kb + && let Some(over_limit_tx) = ceilings.over_limit_tx.take() + && over_limit_tx.send(OverLimit::Disk(disk_kb)).is_err() + { + debug!(disk_kb, "Disk ceiling breached, action already ended"); + } + } + } + // The periodic walk belongs to the soft limit; the guard alone + // measures once, at the end, and pays no walk while running. + samples_since_disk += 1; + if ceilings.disk_limit_kb.is_some() + && samples_since_disk >= DISK_SAMPLE_EVERY + && disk_sample.is_none() + { + samples_since_disk = 0; + disk_sample = start_disk_sample(disk_directory.as_ref()); + } if *stop_rx.borrow() { break; @@ -171,17 +511,35 @@ async fn sample_action_resource_usage( tokio::select! { changed = stop_rx.changed() => { if changed.is_ok() && *stop_rx.borrow() { - sample(&mut peak_memory_kb, &mut cpu_ticks_by_pid); + if let Some(memory_kb) = sample(&mut peak_memory_kb, &mut cpu_ticks_by_pid) { + last_memory_kb = memory_kb; + } break; } } () = tokio::time::sleep(RESOURCE_USAGE_SAMPLE_INTERVAL) => {} } } + // What the action left on disk at the end is the figure that matters + // for sizing; this walk, one stat per entry under the action directory, + // is the only one the guard mode pays. + drop(disk_sample.take()); + if let Some(handle) = start_disk_sample(disk_directory.as_ref()) + && let Ok(disk_kb) = handle.await + { + peak_disk_kb = peak_disk_kb.max(disk_kb); + } SampledResourceUsage { peak_memory_kb, - cpu_time_ms: ticks_to_millis(cpu_ticks_by_pid.values().sum()), + cpu_time_ms: ticks_to_millis( + cpu_ticks_by_pid + .values() + .sum::() + .saturating_sub(baseline_ticks.unwrap_or(0)), + ), + last_memory_kb, + peak_disk_kb, } } @@ -233,7 +591,7 @@ fn sample_process_group(pgid: u32, cpu_ticks_by_pid: &mut HashMap) -> let entry = cpu_ticks_by_pid.entry(member_pid).or_insert(0); *entry = (*entry).max(ticks); } - if let Some(memory_kb) = read_process_rss_kb(member_pid) { + if let Some(memory_kb) = read_process_memory_kb(member_pid) { total_kb += memory_kb; found_any_process = true; } @@ -242,11 +600,27 @@ fn sample_process_group(pgid: u32, cpu_ticks_by_pid: &mut HashMap) -> found_any_process.then_some(total_kb) } +/// The process's share of its resident pages: `Pss` from `smaps_rollup` +/// divides a page shared by N processes N ways, so summed over the group it +/// counts a forked parent's pages, shared libraries and mapped inputs once. +/// `VmRSS` counts them once per process, and a test runner that forks eight +/// children would sample at several times what the cgroup charges. Falls +/// back to `VmRSS` where `smaps_rollup` cannot be read (kernels before 4.14). #[cfg(target_os = "linux")] -fn read_process_rss_kb(pid: u32) -> Option { +fn read_process_memory_kb(pid: u32) -> Option { + if let Ok(rollup) = std::fs::read_to_string(format!("/proc/{pid}/smaps_rollup")) + && let Some(pss_kb) = parse_kb_field(&rollup, "Pss:") + { + return Some(pss_kb); + } let status = std::fs::read_to_string(format!("/proc/{pid}/status")).ok()?; - status.lines().find_map(|line| { - let rest = line.strip_prefix("VmRSS:")?; + parse_kb_field(&status, "VmRSS:") +} + +/// The value of a `Field: 1234 kB` line in `/proc` text, in KiB. +pub fn parse_kb_field(text: &str, field: &str) -> Option { + text.lines().find_map(|line| { + let rest = line.strip_prefix(field)?; rest.split_whitespace().next()?.parse().ok() }) } @@ -1055,7 +1429,14 @@ async fn upload_file( hasher: DigestHasherFunc, metadata: std::fs::Metadata, digest_uploaders: Arc>>, + uploads: Arc, ) -> Result { + // Held from the open to the end of the upload: this is the bound on + // files open at once for the whole output tree. + let _permit = uploads + .acquire() + .await + .err_tip(|| "Upload semaphore closed")?; let is_executable = is_executable(&metadata, &full_path); let file_size = metadata.len(); let file = fs::open_file(&full_path, 0, u64::MAX) @@ -1229,11 +1610,12 @@ fn upload_directory<'a, P: AsRef + Debug + Send + Sync + Clone + 'a>( full_work_directory: &'a str, hasher: DigestHasherFunc, digest_uploaders: Arc>>, + uploads: Arc, ) -> BoxFuture<'a, Result<(Directory, VecDeque), Error>> { Box::pin(async move { - let file_futures = FuturesUnordered::new(); - let dir_futures = FuturesUnordered::new(); - let symlink_futures = FuturesUnordered::new(); + let mut file_futures = Vec::new(); + let mut dir_futures = Vec::new(); + let mut symlink_futures = Vec::new(); { let (_permit, dir_handle) = fs::read_dir(&full_dir_path) .await @@ -1260,6 +1642,7 @@ fn upload_directory<'a, P: AsRef + Debug + Send + Sync + Clone + 'a>( full_work_directory, hasher, digest_uploaders.clone(), + uploads.clone(), ) .and_then(|(dir, all_dirs)| async move { let directory_name = full_path @@ -1294,13 +1677,21 @@ fn upload_directory<'a, P: AsRef + Debug + Send + Sync + Clone + 'a>( ); } else if file_type.is_file() { let digest_uploaders = digest_uploaders.clone(); + let uploads = uploads.clone(); file_futures.push(async move { let metadata = fs::metadata(&full_path) .await .err_tip(|| format!("Could not open file {}", full_path.display()))?; - upload_file(cas_store, &full_path, hasher, metadata, digest_uploaders) - .map_ok(TryInto::try_into) - .await? + upload_file( + cas_store, + &full_path, + hasher, + metadata, + digest_uploaders, + uploads, + ) + .map_ok(TryInto::try_into) + .await? }); } else if file_type.is_symlink() { symlink_futures.push( @@ -1311,12 +1702,20 @@ fn upload_directory<'a, P: AsRef + Debug + Send + Sync + Clone + 'a>( } } - let dir_entries = dir_futures + // The stream bound keeps the fan-out in check at each level; the + // files open at once are bounded by `uploads`, which every level + // shares, since per-level bounds would multiply. + let dir_entries = futures::stream::iter(dir_futures) + .buffer_unordered(UPLOAD_CONCURRENCY) .try_collect::)>>() .await?; let (mut file_nodes, mut symlinks) = try_join( - file_futures.try_collect::>(), - symlink_futures.try_collect::>(), + futures::stream::iter(file_futures) + .buffer_unordered(UPLOAD_CONCURRENCY) + .try_collect::>(), + futures::stream::iter(symlink_futures) + .buffer_unordered(UPLOAD_CONCURRENCY) + .try_collect::>(), ) .await?; @@ -1397,12 +1796,14 @@ async fn do_cleanup( operation_id: &OperationId, action_directory: &str, ) -> Result<(), Error> { - // Mark this operation as being cleaned up - let Some(_cleaning_guard) = running_actions_manager.perform_cleanup(operation_id.clone()) - else { - // Cleanup is already happening elsewhere. - return Ok(()); - }; + // Mark this operation as being cleaned up. The other holders of the + // mark are the orphan sweep and a retry's stale-directory removal, both + // brief; waiting on them is what keeps this cleanup from being skipped, + // which would leave the action's entries and reservations in place for + // good. + let _cleaning_guard = running_actions_manager + .take_cleanup_mark(operation_id.clone()) + .await; let capture = running_actions_manager .buck2_captures @@ -1426,9 +1827,23 @@ async fn do_cleanup( debug!(%action_directory, "Worker cleaning up"); // Note: We need to be careful to keep trying to cleanup even if one of the steps fails. - let remove_dir_result = fs::remove_dir_all(action_directory) - .await - .err_tip(|| format!("Could not remove working directory {action_directory}")); + // A fetch or an upload ended by a kill or a timeout can leave a blocking + // copy or a stamp finishing inside the tree for a moment; the removal is + // retried through that instead of failing on the first try. + let remove_deadline = Instant::now() + running_actions_manager.max_cleanup_wait; + let remove_dir_result = loop { + match fs::remove_dir_all(action_directory).await { + Ok(()) => break Ok(()), + Err(err) if Instant::now() < remove_deadline => { + debug!(%operation_id, ?err, "Removing the working directory again"); + tokio::time::sleep(Duration::from_millis(100)).await; + } + Err(err) => { + break Err(err) + .err_tip(|| format!("Could not remove working directory {action_directory}")); + } + } + }; if let Err(err) = running_actions_manager.cleanup_action(operation_id) { error!(%operation_id, ?err, "Error cleaning up action"); @@ -1471,10 +1886,178 @@ pub trait RunningAction: Sync + Send + Sized + Unpin + 'static { fn get_work_directory(&self) -> &String; } +/// What an action wrote to stdout or stderr: in memory when small, on disk +/// past `max_captured_output_bytes`, with the head kept for the log. +#[derive(Debug)] +pub enum CapturedOutput { + Memory(Bytes), + Spilled { + path: PathBuf, + len: u64, + head: Bytes, + }, +} + +impl CapturedOutput { + const fn len(&self) -> u64 { + match self { + Self::Memory(bytes) => bytes.len() as u64, + Self::Spilled { len, .. } => *len, + } + } + + /// The first bytes, for a log line; all of them when in memory. + fn head(&self) -> &[u8] { + match self { + Self::Memory(bytes) => bytes, + Self::Spilled { head, .. } => head, + } + } +} + +/// The first `limit` bytes of captured output as text for a log line. The +/// cut is by bytes, so one through a multibyte character becomes U+FFFD +/// rather than a panic. +pub fn log_excerpt(bytes: &[u8], limit: usize) -> Cow<'_, str> { + String::from_utf8_lossy(&bytes[..min(bytes.len(), limit)]) +} + +impl From for CapturedOutput { + fn from(bytes: Bytes) -> Self { + Self::Memory(bytes) + } +} + +/// Reads a child's stream to the end, or until `stop` says to keep what +/// has been read: an action can leave a process behind that holds the +/// pipe open after the action itself has exited, and its output must not +/// wait on that. Up to `cap` bytes stay in memory; the moment the buffer +/// passes the cap, everything so far and everything after goes to +/// `spill_path`, and the buffer is cut back to the cap as the head. +pub async fn capture_output( + mut reader: R, + cap: u64, + spill_path: PathBuf, + what: &str, + mut stop: watch::Receiver, +) -> Result { + // `None` when told to stop; a closed sender counts as that too. + async fn read_or_stop( + reader: &mut R, + buf: &mut BytesMut, + stop: &mut watch::Receiver, + what: &str, + ) -> Result, Error> { + if *stop.borrow() { + return Ok(None); + } + tokio::select! { + // Data already in the pipe is read before a stop is honoured. + biased; + read = reader.read_buf(buf) => read + .map(Some) + .err_tip(|| format!("Error reading {what} stream")), + _ = stop.changed() => Ok(None), + } + } + let mut buf = BytesMut::new(); + // A raw handle, outside the open-file budget on purpose: the spill + // lives as long as the child's stream, an fd the budget never covered, + // and a permit held for a whole noisy action would starve the store + // paths the budget exists for. + let mut spill: Option<(tokio::fs::File, u64)> = None; + loop { + if let Some((file, len)) = spill.as_mut() { + let mut chunk = BytesMut::with_capacity(64 * 1024); + let Some(sz) = read_or_stop(&mut reader, &mut chunk, &mut stop, what).await? else { + break; + }; + if sz == 0 { + break; + } + file.write_all(&chunk) + .await + .err_tip(|| format!("Error spilling {what} to {}", spill_path.display()))?; + *len += sz as u64; + continue; + } + let Some(sz) = read_or_stop(&mut reader, &mut buf, &mut stop, what).await? else { + break; + }; + if sz == 0 { + break; + } + if cap > 0 && buf.len() as u64 > cap { + let mut file = tokio::fs::File::create(&spill_path) + .await + .err_tip(|| format!("Error creating {what} spill file {}", spill_path.display()))?; + file.write_all(&buf) + .await + .err_tip(|| format!("Error spilling {what} to {}", spill_path.display()))?; + let len = buf.len() as u64; + buf.truncate(usize::try_from(cap).unwrap_or(usize::MAX)); + spill = Some((file, len)); + } + } + match spill { + None => Ok(CapturedOutput::Memory(buf.freeze())), + Some((mut file, len)) => { + file.flush() + .await + .err_tip(|| format!("Error flushing {what} spill file"))?; + drop(file); + Ok(CapturedOutput::Spilled { + path: spill_path, + len, + head: buf.freeze(), + }) + } + } +} + +/// Uploads captured output to the CAS, from memory or from its spill file. +async fn upload_captured_output( + cas_store: Pin<&FastSlowStore>, + output: CapturedOutput, + hasher: DigestHasherFunc, + what: &str, +) -> Result { + match output { + CapturedOutput::Memory(data) => { + let digest = compute_buf_digest(&data, &mut hasher.hasher()); + cas_store + .update_oneshot(digest, data) + .await + .err_tip(|| format!("Uploading {what}"))?; + Ok(digest) + } + CapturedOutput::Spilled { path, len, .. } => { + let file = fs::open_file(&path, 0, u64::MAX) + .await + .err_tip(|| format!("Could not open {what} spill file {}", path.display()))?; + let (digest, file) = hasher + .hasher() + .digest_for_file(&path, file.into_inner(), Some(len)) + .await + .err_tip(|| format!("Failed to hash {what} spill file {}", path.display()))?; + cas_store + .update_with_whole_file( + digest, + path.clone().into_os_string(), + file, + UploadSizeInfo::ExactSize(len), + ) + .await + .err_tip(|| format!("Uploading {what} from {}", path.display()))?; + Ok(digest) + } + } +} + #[derive(Debug)] struct RunningActionImplExecutionResult { - stdout: Bytes, - stderr: Bytes, + stdout: CapturedOutput, + stderr: CapturedOutput, exit_code: i32, resource_usage: Option, } @@ -1509,6 +2092,9 @@ pub struct RunningActionImpl { input_lease: Option>, timeout: Duration, running_actions_manager: Arc, + /// Set on a kill; every phase that can wait watches it, so a kill ends + /// a fetch or an upload as surely as it ends the command. + kill_token: watch::Sender, state: Mutex, has_manager_entry: AtomicBool, did_cleanup: AtomicBool, @@ -1532,7 +2118,9 @@ impl RunningActionImpl { ) }); let (kill_channel_tx, kill_channel_rx) = oneshot::channel(); + let (kill_token, _) = watch::channel(false); Self { + kill_token, operation_id, action_directory, work_directory, @@ -1575,6 +2163,44 @@ impl RunningActionImpl { /// This function will aggressively download and spawn potentially thousands of futures. It is /// up to the stores to rate limit if needed. async fn inner_prepare_action(self: Arc) -> Result, Error> { + // Before anything lands on disk: an action that reserved more than + // is free is refused as backpressure, not run into ENOSPC halfway + // through its inputs. Free means free after what the actions + // already admitted reserved: four 4 GiB actions dispatched in one + // second all see the same 10 GiB, so the disk alone would admit + // them all. The reservation is held until the action's cleanup. + if let Some(property) = self + .running_actions_manager + .execution_configuration + .resource_enforcement + .as_ref() + .and_then(|enforcement| enforcement.disk_property_name.as_deref()) + && let Some(reserved_kb) = self + .action_info + .platform_properties + .get(property) + .and_then(|value| value.parse::().ok()) + && reserved_kb > 0 + && let Some(free_kb) = free_disk_kb(&self.running_actions_manager.root_action_directory) + { + // Held across the check and the insert, so two actions + // preparing at once cannot both pass on the same free space. + let mut reservations = self.running_actions_manager.disk_reservations_kb.lock(); + let reserved_by_others_kb: u64 = reservations + .iter() + .filter(|(operation_id, _)| **operation_id != self.operation_id) + .map(|(_, kb)| *kb) + .sum(); + if free_kb.saturating_sub(reserved_by_others_kb) < reserved_kb { + self.metrics().disk_guard_refusals.inc(); + return Err(make_err!( + Code::ResourceExhausted, + "Not enough free disk for this action: reserved {reserved_kb} KiB ({property}), {free_kb} KiB free under {} with {reserved_by_others_kb} KiB reserved by the actions already admitted", + self.running_actions_manager.root_action_directory + )); + } + reservations.insert(self.operation_id.clone(), reserved_kb); + } { let mut state = self.state.lock(); state.execution_metadata.input_fetch_start_timestamp = @@ -1589,6 +2215,7 @@ impl RunningActionImpl { ) .await .err_tip(|| "Converting command_digest to Command") + .map_err(tag_missing_input) }); let filesystem_store_pin = Pin::new(self.running_actions_manager.filesystem_store.as_ref()); @@ -1611,6 +2238,7 @@ impl RunningActionImpl { self.input_lease.clone(), )) .await + .map_err(tag_missing_input) }) .await?; command @@ -1689,6 +2317,98 @@ impl RunningActionImpl { }) } + /// The environment an action's process starts with: the worker's + /// `additional_environment`, then the action's own variables on top, + /// and on Windows the `SystemRoot` and `PATH` defaults without which + /// nothing runs. With `per_action` the sources that belong to one + /// action (its timeout, a side channel file, its directory) are in and + /// the side channel file is reported through it; without it they are + /// left out, for a pooled process that serves many actions. + fn action_environment( + &self, + command_proto: &ProtoCommand, + requested_timeout: Duration, + mut per_action: Option<&mut Option>>, + ) -> BTreeMap { + let mut env = BTreeMap::new(); + if let Some(additional_environment) = &self + .running_actions_manager + .execution_configuration + .additional_environment + { + for (name, source) in additional_environment { + let value = match source { + EnvironmentSource::Property(property) => self + .action_info + .platform_properties + .get(property) + .cloned() + .unwrap_or_default(), + EnvironmentSource::Value(value) => value.clone(), + EnvironmentSource::FromEnvironment => env::var(name).unwrap_or_default(), + EnvironmentSource::TimeoutMillis => { + if per_action.is_none() { + continue; + } + requested_timeout.as_millis().to_string() + } + EnvironmentSource::SideChannelFile => { + let Some(side_channel_file) = per_action.as_mut() else { + continue; + }; + let file = format!("{}/{}", self.action_directory, Uuid::new_v4().simple()); + **side_channel_file = Some(Cow::Owned(file.clone().into())); + file + } + EnvironmentSource::ActionDirectory => { + if per_action.is_none() { + continue; + } + self.action_directory.clone() + } + }; + env.insert(name.clone(), value); + } + } + for variable in &command_proto.environment_variables { + env.insert(variable.name.clone(), variable.value.clone()); + } + // If SystemRoot is not set on windows we set it to default. Failing + // to do this causes all commands to fail. + #[cfg(target_family = "windows")] + { + if !env.keys().any(|name| name.to_uppercase() == "SYSTEMROOT") { + env.insert("SystemRoot".to_string(), "C:\\Windows".to_string()); + } + if !env.keys().any(|name| name.to_uppercase() == "PATH") { + env.insert("PATH".to_string(), "C:\\Windows\\System32".to_string()); + } + } + env + } + + /// What the scheduler's kill leaves behind for a pooled action, whether + /// it came while waiting for a worker process or while one served it. + fn record_persistent_worker_kill(&self, command_proto: ProtoCommand, command: &str) { + let mut state = self.state.lock(); + state.error = Error::merge_option( + state.error.take(), + Some(Error::new( + Code::Cancelled, + format!("Persistent worker command '{command}' was killed by scheduler"), + )), + ); + state.command_proto = Some(command_proto); + state.execution_result = Some(RunningActionImplExecutionResult { + stdout: Bytes::new().into(), + stderr: Bytes::new().into(), + exit_code: EXIT_CODE_FOR_SIGNAL, + resource_usage: None, + }); + state.execution_metadata.execution_completed_timestamp = + (self.running_actions_manager.callbacks.now_fn)(); + } + async fn inner_execute(self: Arc) -> Result, Error> { let (command_proto, mut kill_channel_rx) = { let mut state = self.state.lock(); @@ -1739,11 +2459,28 @@ impl RunningActionImpl { let program = self .canonicalise_path(args[0], &command_proto.working_directory) .err_tip(|| format!("Canonicalisation failure. Command={args:#?}"))?; - if let Some(wire_format_result) = action_supports_persistent_workers(&self.action_info) { + let requested_timeout = if self.action_info.timeout.is_zero() { + self.running_actions_manager.max_action_timeout + } else { + self.action_info.timeout + }; + if self + .running_actions_manager + .execution_configuration + .persistent_workers + .enabled + && let Some(wire_format_result) = action_supports_persistent_workers(&self.action_info) + { match wire_format_result { Ok(wire_format) => { let command_argv = os_args_to_strings(&args)?; - let key = WorkerKey::from_argv(&command_argv, wire_format)?; + // What a one-shot action gets, but for the sources that + // belong to one action: a process that serves many + // cannot carry a timeout, a side channel file or an + // action directory. + let env = self.action_environment(&command_proto, requested_timeout, None); + let key = WorkerKey::from_argv(&command_argv, wire_format)? + .with_env(env.into_iter().collect()); let request = WorkRequest { arguments: persistent_worker_request_arguments(&command_argv), inputs: Vec::::new(), @@ -1755,23 +2492,81 @@ impl RunningActionImpl { let worker_cwd = PathBuf::from(&self.running_actions_manager.root_action_directory); - match self + // The wait for a process is bounded by the action's own + // timeout as well as the pool's, and raced against the + // kill: a cancelled action must not hold a slot, and a + // short action cannot wait longer than it may run. + let acquire_started = Instant::now(); + let acquire_bound = self .running_actions_manager .persistent_worker_pool - .acquire(key.clone(), &program, &worker_cwd) - .await - { + .acquire_timeout() + .min(self.timeout); + let acquired = { + let acquire_fut = self + .running_actions_manager + .persistent_worker_pool + .acquire(key.clone(), &program, &worker_cwd); + tokio::pin!(acquire_fut); + tokio::select! { + result = &mut acquire_fut => Some(result), + _ = &mut kill_channel_rx => None, + () = tokio::time::sleep(acquire_bound) => Some(Err(make_err!( + Code::DeadlineExceeded, + "No persistent worker came back within the action's remaining time ({acquire_bound:?})" + ))), + } + }; + let Some(acquired) = acquired else { + let command = args.join(OsStr::new(" ")).to_string_lossy().into_owned(); + self.record_persistent_worker_kill(command_proto, &command); + return Ok(self); + }; + match acquired { Ok(mut lease) => { + self.metrics().persistent_worker_dispatches.inc(); let timer = self.metrics().child_process.begin_timer(); + let execution_started = Instant::now(); + // What the pooled process and its children use + // while they serve this request. The process is + // shared, so nothing here is enforced against + // the action's reservation; it is measured so + // the sizing loop learns what the tool needs. + // CPU is what the process spent from here on; + // peak memory is the whole process, which holds + // what earlier requests left in it. + #[cfg(target_os = "linux")] + let sampler = lease.worker().pid().map(|pgid| { + start_action_resource_usage_sampler( + pgid, + Ceilings { + memory_limit_kb: None, + disk_limit_kb: None, + over_limit_tx: None, + }, + None, + ) + }); let dispatch_result = { + let remaining = + self.timeout.saturating_sub(acquire_started.elapsed()); let dispatch_fut = - lease.worker().dispatch_with_timeout(&request, self.timeout); + lease.worker().dispatch_with_timeout(&request, remaining); tokio::pin!(dispatch_fut); tokio::select! { result = &mut dispatch_fut => Some(result), _ = &mut kill_channel_rx => None, } }; + #[cfg(target_os = "linux")] + let sampled_usage = match sampler { + Some(sampler) => finish_action_resource_usage_sampler(sampler) + .await + .unwrap_or_default(), + None => SampledResourceUsage::default(), + }; + #[cfg(not(target_os = "linux"))] + let sampled_usage = SampledResourceUsage::default(); let response = match dispatch_result { Some(Ok(response)) => { lease.release(true).await; @@ -1786,29 +2581,9 @@ impl RunningActionImpl { None => { drop(timer); lease.release(false).await; - { - let mut state = self.state.lock(); - state.error = Error::merge_option( - state.error.take(), - Some(Error::new( - Code::Cancelled, - format!( - "Persistent worker command '{}' was killed by scheduler", - args.join(OsStr::new(" ")).to_string_lossy() - ), - )), - ); - state.command_proto = Some(command_proto); - state.execution_result = - Some(RunningActionImplExecutionResult { - stdout: Bytes::new(), - stderr: Bytes::new(), - exit_code: EXIT_CODE_FOR_SIGNAL, - resource_usage: None, - }); - state.execution_metadata.execution_completed_timestamp = - (self.running_actions_manager.callbacks.now_fn)(); - } + let command = + args.join(OsStr::new(" ")).to_string_lossy().into_owned(); + self.record_persistent_worker_kill(command_proto, &command); return Ok(self); } }; @@ -1820,14 +2595,47 @@ impl RunningActionImpl { self.metrics().child_process_failure_error_code.inc(); } info!(?args, ?key, "Persistent worker command complete"); + let (memory_property, disk_property) = self + .running_actions_manager + .execution_configuration + .resource_enforcement + .as_ref() + .map_or((None, None), |e| { + ( + e.memory.as_ref().map(|m| m.property_name.as_str()), + e.disk_property_name.as_deref(), + ) + }); + let (outcome, enforced) = + classify_outcome(None, None, sampled_usage.last_memory_kb, None); + let resource_usage = Some(ActionResourceUsage { + peak_memory_kb: sampled_usage.peak_memory_kb, + cpu_time_ms: sampled_usage.cpu_time_ms, + sampled: sampled_usage.peak_memory_kb > 0 + || sampled_usage.cpu_time_ms > 0, + operation_id: String::new(), + worker_id: String::new(), + wall_time_ms: u64::try_from( + execution_started.elapsed().as_millis(), + ) + .unwrap_or(u64::MAX), + peak_disk_kb: 0, + outcome: outcome.into(), + enforced, + reserved: reservation_of( + &self.action_info.platform_properties, + memory_property, + disk_property, + ), + }); { let mut state = self.state.lock(); state.command_proto = Some(command_proto); state.execution_result = Some(RunningActionImplExecutionResult { - stdout: Bytes::new(), - stderr: Bytes::from(response.output), + stdout: Bytes::new().into(), + stderr: Bytes::from(response.output).into(), exit_code: response.exit_code, - resource_usage: None, + resource_usage, }); state.execution_metadata.execution_completed_timestamp = (self.running_actions_manager.callbacks.now_fn)(); @@ -1835,6 +2643,7 @@ impl RunningActionImpl { return Ok(self); } Err(err) => { + self.metrics().persistent_worker_fallbacks.inc(); info!( ?err, ?key, @@ -1844,6 +2653,7 @@ impl RunningActionImpl { } } Err(err) => { + self.metrics().persistent_worker_fallbacks.inc(); info!( ?err, "Falling back to one-shot execution; unsupported persistent worker protocol" @@ -1864,73 +2674,28 @@ impl RunningActionImpl { .current_dir(command_directory.clone()) .env_clear(); - let requested_timeout = if self.action_info.timeout.is_zero() { - self.running_actions_manager.max_action_timeout - } else { - self.action_info.timeout - }; - - let mut maybe_side_channel_file: Option> = None; - if let Some(additional_environment) = &self + if self .running_actions_manager .execution_configuration - .additional_environment + .set_tmpdir { - for (name, source) in additional_environment { - let value = match source { - EnvironmentSource::Property(property) => self - .action_info - .platform_properties - .get(property) - .map_or_else(|| Cow::Borrowed(""), |v| Cow::Borrowed(v.as_str())), - EnvironmentSource::Value(value) => Cow::Borrowed(value.as_str()), - EnvironmentSource::FromEnvironment => { - Cow::Owned(env::var(name).unwrap_or_default()) - } - EnvironmentSource::TimeoutMillis => { - Cow::Owned(requested_timeout.as_millis().to_string()) - } - EnvironmentSource::SideChannelFile => { - let file_cow = - format!("{}/{}", self.action_directory, Uuid::new_v4().simple()); - maybe_side_channel_file = Some(Cow::Owned(file_cow.clone().into())); - Cow::Owned(file_cow) - } - EnvironmentSource::ActionDirectory => { - Cow::Borrowed(self.action_directory.as_str()) - } - }; - command_builder.env(name, value.as_ref()); - } + // One tmp per action, removed with the action directory, so two + // actions writing the same name under $TMPDIR no longer collide. + // Code that hardcodes /tmp still does. + let tmp_directory = format!("{}/tmp", self.action_directory); + fs::create_dir_all(&tmp_directory) + .await + .err_tip(|| format!("Creating {tmp_directory} for TMPDIR"))?; + command_builder.env("TMPDIR", &tmp_directory); } - #[cfg(target_family = "unix")] - let envs = &command_proto.environment_variables; - // If SystemRoot is not set on windows we set it to default. Failing to do - // this causes all commands to fail. - #[cfg(target_family = "windows")] - let envs = { - let mut envs = command_proto.environment_variables.clone(); - if !envs.iter().any(|v| v.name.to_uppercase() == "SYSTEMROOT") { - envs.push( - nativelink_proto::build::bazel::remote::execution::v2::command::EnvironmentVariable { - name: "SystemRoot".to_string(), - value: "C:\\Windows".to_string(), - }, - ); - } - if !envs.iter().any(|v| v.name.to_uppercase() == "PATH") { - envs.push( - nativelink_proto::build::bazel::remote::execution::v2::command::EnvironmentVariable { - name: "PATH".to_string(), - value: "C:\\Windows\\System32".to_string(), - }, - ); - } - envs - }; - for environment_variable in envs { - command_builder.env(&environment_variable.name, &environment_variable.value); + let mut maybe_side_channel_file: Option> = None; + for (name, value) in self.action_environment( + &command_proto, + requested_timeout, + Some(&mut maybe_side_channel_file), + ) { + command_builder.env(name, value); } // Sandboxing of the command if we are running on Linux, this resolves issues where @@ -1940,42 +2705,64 @@ impl RunningActionImpl { let use_namespaces = self.running_actions_manager.use_namespaces; if !matches!(use_namespaces, UseNamespaces::No) { + let (mount, isolate_tmp) = match use_namespaces { + UseNamespaces::YesAndMount { isolate_tmp } => (true, isolate_tmp), + UseNamespaces::No | UseNamespaces::Yes => (false, false), + }; let root_action_directory = std::ffi::CString::new( self.running_actions_manager.root_action_directory.clone(), ) .err_tip(|| "In RunningActionImpl::inner_execute()")?; let action_directory = std::ffi::CString::new(self.action_directory.clone()) .err_tip(|| "In RunningActionImpl::inner_execute()")?; + // The action's private /tmp is its own tmp directory, the + // same one TMPDIR points at, bound over /tmp in the child. + let tmp_directory = if isolate_tmp { + let tmp_directory = format!("{}/tmp", self.action_directory); + fs::create_dir_all(&tmp_directory) + .await + .err_tip(|| format!("Creating {tmp_directory} for the private /tmp"))?; + Some( + std::ffi::CString::new(tmp_directory) + .err_tip(|| "In RunningActionImpl::inner_execute()")?, + ) + } else { + None + }; // SAFETY: This function is specifically designed to operate in a async-signal-safe // environment. unsafe { command_builder.pre_exec(move || { crate::namespace_utils::configure_namespace( - matches!(use_namespaces, UseNamespaces::YesAndMount), + mount, + tmp_directory.as_deref(), &root_action_directory, &action_directory, ) }); } } - - // Run the action as its own process-group leader (pgid == child - // pid). The resource-usage sampler attributes memory by process - // group, so this keeps the whole action together — including - // processes reparented to the worker when an intermediate shell - // exits — and never conflates it with the worker's own group. - command_builder.process_group(0); } + // Run the action as its own process-group leader (pgid == child + // pid) on every unix. The resource-usage sampler attributes memory + // by process group, and the kill at exit ends the group, so this + // keeps the whole action together, including processes reparented + // to the worker when an intermediate shell exits, and never + // conflates it with the worker's own group. + #[cfg(target_family = "unix")] + command_builder.process_group(0); let mut child_process = command_builder .spawn() .err_tip(|| format!("Could not execute command {args:?}"))?; - let mut stdout_reader = child_process + // Ours until the wait below collects it; the reaper leaves it alone. + let _owned_child = crate::reaper::OwnedChild::new(child_process.id()); + let stdout_reader = child_process .stdout .take() .err_tip(|| "Expected stdout to exist on command this should never happen")?; - let mut stderr_reader = child_process + let stderr_reader = child_process .stderr .take() .err_tip(|| "Expected stderr to exist on command this should never happen")?; @@ -1991,10 +2778,83 @@ impl RunningActionImpl { child_process, ); + // The reservation the scheduler placed this action by, if the + // worker is asked to hold it to that number. + let memory_reservation = self + .running_actions_manager + .execution_configuration + .resource_enforcement + .as_ref() + .and_then(|enforcement| { + let memory = enforcement.memory.as_ref()?; + let reserved_kb = self + .action_info + .platform_properties + .get(&memory.property_name)? + .parse::() + .ok()?; + (reserved_kb > 0).then(|| { + let limit_kb = reserved_kb.saturating_mul(100 + memory.headroom_percent) / 100; + (reserved_kb, limit_kb) + }) + }); + // The disk reservation, held to only with `disk: soft`; disk is + // measured whenever the enforcement names a disk property. + let disk_enforcement = self + .running_actions_manager + .execution_configuration + .resource_enforcement + .as_ref() + .and_then(|enforcement| { + let property = enforcement.disk_property_name.as_deref()?; + let reserved_kb = self + .action_info + .platform_properties + .get(property) + .and_then(|value| value.parse::().ok()) + .unwrap_or(0); + let limit_kb = (enforcement.disk_soft && reserved_kb > 0).then(|| { + reserved_kb.saturating_mul(100 + enforcement.disk_headroom_percent) / 100 + }); + Some((reserved_kb, limit_kb)) + }); + let (over_limit_tx, over_limit_rx) = oneshot::channel::(); + // Holds the sender when no ceiling is enforced, so the receiver + // below stays pending instead of resolving closed. + let mut over_limit_keepalive = Some(over_limit_tx); #[cfg(target_os = "linux")] - let mut maybe_resource_usage_sampler = - child_process.id().map(start_action_resource_usage_sampler); + let mut maybe_resource_usage_sampler = child_process.id().map(|pgid| { + let memory_limit_kb = memory_reservation.map(|(_, limit_kb)| limit_kb); + let disk_limit_kb = disk_enforcement.and_then(|(_, limit_kb)| limit_kb); + let over_limit_tx = (memory_limit_kb.is_some() || disk_limit_kb.is_some()) + .then(|| over_limit_keepalive.take()) + .flatten(); + let disk_directory = disk_enforcement.map(|_| { + let directory = PathBuf::from(&self.action_directory); + let since = directory_stamp(&directory).unwrap_or_else(|err| { + warn!(?err, "Timing the disk walk by the wall clock instead"); + SystemTime::now() + }); + (directory, since) + }); + start_action_resource_usage_sampler( + pgid, + Ceilings { + memory_limit_kb, + disk_limit_kb, + over_limit_tx, + }, + disk_directory, + ) + }); + #[cfg(not(target_os = "linux"))] + let _ = &disk_enforcement; + let mut over_limit_fut = over_limit_rx.fuse(); + // The group to end when the action's own process has exited; taken + // before the guard owns the child. + #[cfg(target_family = "unix")] + let action_pgid = child_process.id(); let mut child_process_guard = guard(child_process, |mut child_process| { let result: Result, std::io::Error> = child_process.try_wait(); @@ -2014,33 +2874,45 @@ impl RunningActionImpl { } }); + let output_cap = self + .running_actions_manager + .execution_configuration + .max_captured_output_bytes; + let stdout_spill = PathBuf::from(&self.action_directory).join("nativelink.stdout"); + let stderr_spill = PathBuf::from(&self.action_directory).join("nativelink.stderr"); + // Flipped when the pipes have had their drain after the action + // ended; the readers then keep what they have. + let (drain_stop_tx, drain_stop_rx) = watch::channel(false); + let stdout_stop = drain_stop_rx.clone(); let all_stdout_fut = spawn!("stdout_reader", async move { - let mut all_stdout = BytesMut::new(); - loop { - let sz = stdout_reader - .read_buf(&mut all_stdout) - .await - .err_tip(|| "Error reading stdout stream")?; - if sz == 0 { - break; // EOF. - } - } - Result::::Ok(all_stdout.freeze()) + capture_output( + stdout_reader, + output_cap, + stdout_spill, + "stdout", + stdout_stop, + ) + .await }); let all_stderr_fut = spawn!("stderr_reader", async move { - let mut all_stderr = BytesMut::new(); - loop { - let sz = stderr_reader - .read_buf(&mut all_stderr) - .await - .err_tip(|| "Error reading stderr stream")?; - if sz == 0 { - break; // EOF. - } - } - Result::::Ok(all_stderr.freeze()) + capture_output( + stderr_reader, + output_cap, + stderr_spill, + "stderr", + drain_stop_rx, + ) + .await }); - let mut killed_action = false; + let mut kill_reason: Option = None; + let kill_grace = self + .running_actions_manager + .execution_configuration + .kill_grace; + // Armed by a soft kill; fires SIGKILL if the action is still there + // when the grace period ends. + let mut grace_fut: Fuse> = Fuse::terminated(); + let execution_started = std::time::Instant::now(); let timer = self.metrics().child_process.begin_timer(); let mut sleep_fut = (self.running_actions_manager.callbacks.sleep_fn)(self.timeout).fuse(); @@ -2048,13 +2920,8 @@ impl RunningActionImpl { tokio::select! { () = &mut sleep_fut => { self.running_actions_manager.metrics.task_timeouts.inc(); - killed_action = true; - if let Err(err) = child_process_guard.kill().await { - error!( - ?err, - "Could not kill process in RunningActionsManager for action timeout", - ); - } + kill_reason = Some(KillReason::Timeout); + grace_fut = stop_action(&mut child_process_guard, kill_grace, "action timeout").await; { let joined_command = args.join(OsStr::new(" ")); let command = joined_command.to_string_lossy(); @@ -2075,40 +2942,174 @@ impl RunningActionImpl { ))); } }, + () = &mut grace_fut => { + warn!( + operation_id = ?self.operation_id, + grace_ms = kill_grace.as_millis(), + "Action did not leave on SIGTERM within its grace period, sending SIGKILL" + ); + hard_kill(&mut child_process_guard, "grace period expired").await; + }, + Ok(over_limit) = &mut over_limit_fut => match over_limit { + OverLimit::Disk(observed_kb) => { + self.running_actions_manager.metrics.disk_reservation_kills.inc(); + kill_reason = Some(KillReason::Disk); + // The files are still being written; no grace here. + hard_kill(&mut child_process_guard, "disk reservation").await; + let (reserved_kb, limit_kb) = disk_enforcement + .map_or((0, 0), |(reserved_kb, limit_kb)| { + (reserved_kb, limit_kb.unwrap_or(0)) + }); + let (property, headroom) = self + .running_actions_manager + .execution_configuration + .resource_enforcement + .as_ref() + .map_or(("", 0), |e| (e.disk_property_name.as_deref().unwrap_or(""), e.disk_headroom_percent)); + warn!( + operation_id = ?self.operation_id, + reserved_kb, + limit_kb, + observed_kb, + "Action exceeded its disk reservation, killed" + ); + let mut state = self.state.lock(); + state.error = Error::merge_option(state.error.take(), Some(Error::new( + Code::FailedPrecondition, + format!( + "Action exceeded its disk reservation: reserved {reserved_kb} KiB ({property}), \ + limit {limit_kb} KiB with {headroom}% headroom, its own files came to {observed_kb} KiB. \ + Raise the reservation or write less." + ), + ))); + + } + OverLimit::Memory(observed_kb) => { + self.running_actions_manager.metrics.memory_reservation_kills.inc(); + kill_reason = Some(KillReason::Memory); + // Memory is still growing; no grace here. + hard_kill(&mut child_process_guard, "memory reservation").await; + let (reserved_kb, limit_kb) = memory_reservation.unwrap_or((0, 0)); + let (property, headroom) = self + .running_actions_manager + .execution_configuration + .resource_enforcement + .as_ref() + .and_then(|e| e.memory.as_ref()) + .map_or(("", 0), |m| (m.property_name.as_str(), m.headroom_percent)); + warn!( + operation_id = ?self.operation_id, + reserved_kb, + limit_kb, + observed_kb, + "Action exceeded its memory reservation, killed" + ); + let mut state = self.state.lock(); + state.error = Error::merge_option(state.error.take(), Some(Error::new( + Code::FailedPrecondition, + format!( + "Action exceeded its memory reservation: reserved {reserved_kb} KiB ({property}), \ + limit {limit_kb} KiB with {headroom}% headroom, observed {observed_kb} KiB. \ + Raise the reservation or shrink the action." + ), + ))); + + } + }, maybe_exit_status = child_process_guard.wait() => { // Defuse our guard so it does not try to cleanup and make senseless logs. drop(ScopeGuard::<_, _>::into_inner(child_process_guard)); + drop(over_limit_keepalive.take()); let exit_status = maybe_exit_status.err_tip(|| "Failed to collect exit code of process")?; // TODO(palfrey) We should implement stderr/stdout streaming to client here. // If we get killed before the stream is started, then these will lock up. - // TODO(palfrey) There is a significant bug here. If we kill the action and the action creates - // child processes, it can create zombies. See: https://github.com/tracemachina/nativelink/issues/225 - let (stdout, stderr) = if killed_action { + let killed_action = kill_reason.is_some(); + if killed_action { drop(timer); - (Bytes::new(), Bytes::new()) } else { timer.measure(); - let (maybe_all_stdout, maybe_all_stderr) = tokio::join!(all_stdout_fut, all_stderr_fut); - ( - maybe_all_stdout.err_tip(|| "Internal error reading from stdout of worker task")??, - maybe_all_stderr.err_tip(|| "Internal error reading from stderr of worker task")?? - ) + } + // The action is over; anything still running in its + // group is a straggler (a background job, a helper that + // daemonized) holding the pipes open, and without this + // the action would sit here until its timeout. The + // sandboxes Bazel runs locally end them the same way. + #[cfg(target_family = "unix")] + let had_stragglers = action_pgid + .is_some_and(|pgid| signal_process_group(pgid, libc::SIGKILL)); + #[cfg(not(target_family = "unix"))] + let had_stragglers = false; + // An action that exited on its own with nothing left in + // its group closes its pipes itself, and its output is + // the result, so that read runs to the end: a reader + // slowed by a busy spill must not lose bytes to a clock. + // The wait is bounded only when there was something to + // hold the pipes: a killed action, whose output is a + // bonus, or stragglers, one of which may have left the + // group. What was read by then is kept. + let mut drain = core::pin::pin!(async { tokio::join!(all_stdout_fut, all_stderr_fut) }); + let (maybe_all_stdout, maybe_all_stderr) = if killed_action || had_stragglers { + tokio::select! { + outputs = &mut drain => outputs, + () = tokio::time::sleep(OUTPUT_DRAIN_TIMEOUT) => { + warn!( + operation_id = ?self.operation_id, + drain_s = OUTPUT_DRAIN_TIMEOUT.as_secs(), + killed_action, + had_stragglers, + "A process held the action's output pipes open past the drain; keeping what was read" + ); + drain_stop_tx.send_replace(true); + drain.await + } + } + } else { + drain.await + }; + let (stdout, stderr) = ( + maybe_all_stdout.err_tip(|| "Internal error reading from stdout of worker task")??, + maybe_all_stderr.err_tip(|| "Internal error reading from stderr of worker task")??, + ); + + #[cfg(target_os = "linux")] + let sampled_usage = match maybe_resource_usage_sampler.take() { + Some(sampler) => finish_action_resource_usage_sampler(sampler) + .await + .unwrap_or_default(), + None => SampledResourceUsage::default(), }; + #[cfg(not(target_os = "linux"))] + let sampled_usage = SampledResourceUsage::default(); + + let memory_limit_kb = memory_reservation + .map(|(_, limit_kb)| limit_kb) + .or(self.running_actions_manager.pod_memory_limit_kb); + // The signal that ended the process, where the platform reports one. + #[cfg(target_family = "unix")] + let signal = exit_status.signal(); + #[cfg(not(target_family = "unix"))] + let signal = None; + let (outcome, enforced) = classify_outcome( + kill_reason, + signal, + sampled_usage.last_memory_kb, + memory_limit_kb, + ); let exit_code = exit_status.code().map_or_else(|| { - // No exit code means the runner was terminated by a - // signal. SIGKILL on Linux is the kernel OOM killer's - // weapon of choice, so flag this for operators trying - // to correlate action failures with kubectl-top - // memory pressure. - warn!( - ?args, - "Runner subprocess terminated by signal (no exit code); likely OOMKilled \ - or externally killed. If this repeats for the same action, raise \ - `workers.specs[*].resources.limits.memory` or shrink the action's \ - concurrency." - ); self.metrics().child_process_failure_error_code.inc(); + if kill_reason.is_none() { + // No exit code and no kill of ours: a signal from + // outside. SIGKILL near the memory limit is the + // kernel's OOM killer. + warn!( + ?args, + outcome = outcome.as_str_name(), + last_memory_kb = sampled_usage.last_memory_kb, + ?memory_limit_kb, + "Runner subprocess terminated by a signal nobody here sent" + ); + } EXIT_CODE_FOR_SIGNAL }, |exit_code| { if exit_code == 0 { @@ -2118,32 +3119,53 @@ impl RunningActionImpl { } exit_code }); + if kill_reason.is_none() && outcome == ResourceOutcome::KilledMemory { + let mut state = self.state.lock(); + state.error = Error::merge_option(state.error.take(), Some(Error::new( + Code::FailedPrecondition, + format!( + "Action was killed by the kernel at {} KiB, within 10% of the {} KiB memory limit: \ + out of memory. Raise the reservation or shrink the action.", + sampled_usage.last_memory_kb, + memory_limit_kb.unwrap_or(0), + ), + ))); + } - #[cfg(target_os = "linux")] - let resource_usage = match maybe_resource_usage_sampler.take() { - Some(sampler) => finish_action_resource_usage_sampler(sampler) - .await - .and_then(|usage| { - // An action too short to catch a sample leaves - // both at zero; report nothing rather than a - // misleading zero. - (usage.peak_memory_kb > 0 || usage.cpu_time_ms > 0).then_some( - ActionResourceUsage { - peak_memory_kb: usage.peak_memory_kb, - cpu_time_ms: usage.cpu_time_ms, - sampled: true, - operation_id: String::new(), - worker_id: String::new(), - }, - ) - }), - None => None, - }; - #[cfg(not(target_os = "linux"))] - let resource_usage = None; + let (memory_property, disk_property) = self + .running_actions_manager + .execution_configuration + .resource_enforcement + .as_ref() + .map_or((None, None), |e| { + ( + e.memory.as_ref().map(|m| m.property_name.as_str()), + e.disk_property_name.as_deref(), + ) + }); + let resource_usage = Some(ActionResourceUsage { + peak_memory_kb: sampled_usage.peak_memory_kb, + cpu_time_ms: sampled_usage.cpu_time_ms, + // An action too short to catch a sample leaves both at + // zero; say so rather than report a misleading zero. + sampled: sampled_usage.peak_memory_kb > 0 + || sampled_usage.cpu_time_ms > 0 + || sampled_usage.peak_disk_kb > 0, + operation_id: String::new(), + worker_id: String::new(), + wall_time_ms: u64::try_from(execution_started.elapsed().as_millis()).unwrap_or(u64::MAX), + peak_disk_kb: sampled_usage.peak_disk_kb, + outcome: outcome.into(), + enforced, + reserved: reservation_of( + &self.action_info.platform_properties, + memory_property, + disk_property, + ), + }); // log something useful instead of repeating same ?arg - info!(?exit_code, "Command complete"); + info!(?exit_code, outcome = outcome.as_str_name(), "Command complete"); let maybe_error_override = if let Some(side_channel_file) = maybe_side_channel_file { process_side_channel_file(side_channel_file.clone(), &args, requested_timeout).await @@ -2167,14 +3189,8 @@ impl RunningActionImpl { return Ok(self); }, _ = &mut kill_channel_rx => { - killed_action = true; - if let Err(err) = child_process_guard.kill().await { - error!( - operation_id = ?self.operation_id, - ?err, - "Could not kill process", - ); - } + kill_reason = Some(KillReason::External); + grace_fut = stop_action(&mut child_process_guard, kill_grace, "scheduler kill").await; { let mut state = self.state.lock(); state.error = Error::merge_option(state.error.take(), Some(Error::new( @@ -2224,7 +3240,7 @@ impl RunningActionImpl { let cas_store = self.running_actions_manager.cas_store.as_ref(); let hasher = self.action_info.unique_qualifier.digest_function(); - let mut output_path_futures = FuturesUnordered::new(); + let mut output_path_futures: Vec>> = Vec::new(); let mut output_paths = command_proto.output_paths; if output_paths.is_empty() { output_paths @@ -2233,6 +3249,7 @@ impl RunningActionImpl { output_paths.append(&mut command_proto.output_directories); } let digest_uploaders = Arc::new(Mutex::new(HashMap::new())); + let uploads = Arc::new(Semaphore::new(UPLOAD_CONCURRENCY)); for entry in output_paths { let full_path = OsString::from(if command_proto.working_directory.is_empty() { format!("{}/{}", self.work_directory, entry) @@ -2244,198 +3261,210 @@ impl RunningActionImpl { }); let work_directory = &self.work_directory; let digest_uploaders = digest_uploaders.clone(); - output_path_futures.push(async move { - let metadata = { - let metadata = match fs::symlink_metadata(&full_path).await { - Ok(file) => file, - Err(e) => { - if e.code == Code::NotFound { - // In the event our output does not exist, according to the bazel remote - // execution spec, we simply ignore it continue. - return Result::::Ok(OutputType::None); + let uploads = uploads.clone(); + output_path_futures.push( + async move { + let metadata = { + let metadata = match fs::symlink_metadata(&full_path).await { + Ok(file) => file, + Err(e) => { + if e.code == Code::NotFound { + // In the event our output does not exist, according to the bazel remote + // execution spec, we simply ignore it continue. + return Result::::Ok(OutputType::None); + } + return Err(e).err_tip(|| { + format!("Could not open file {}", full_path.display()) + }); } - return Err(e).err_tip(|| { - format!("Could not open file {}", full_path.display()) - }); + }; + + if metadata.is_file() { + return Ok(OutputType::File( + upload_file( + cas_store.as_pin(), + &full_path, + hasher, + metadata, + digest_uploaders, + uploads, + ) + .await + .map(|mut file_info| { + file_info.name_or_path = NameOrPath::Path(entry); + file_info + }) + .err_tip(|| format!("Uploading file {}", full_path.display()))?, + )); } + metadata }; - - if metadata.is_file() { - return Ok(OutputType::File( - upload_file( + if metadata.is_dir() { + Ok(OutputType::Directory( + upload_directory( cas_store.as_pin(), &full_path, + work_directory, hasher, - metadata, digest_uploaders, + uploads, ) - .await - .map(|mut file_info| { - file_info.name_or_path = NameOrPath::Path(entry); - file_info + .and_then(|(root_dir, children)| async move { + let tree = ProtoTree { + root: Some(root_dir), + children: children.into(), + }; + let tree_digest = serialize_and_upload_message( + &tree, + cas_store.as_pin(), + &mut hasher.hasher(), + ) + .await + .err_tip(|| format!("While processing {entry}"))?; + Ok(DirectoryInfo { + path: entry, + tree_digest, + }) }) - .err_tip(|| format!("Uploading file {}", full_path.display()))?, - )); - } - metadata - }; - if metadata.is_dir() { - Ok(OutputType::Directory( - upload_directory( - cas_store.as_pin(), - &full_path, - work_directory, - hasher, - digest_uploaders, - ) - .and_then(|(root_dir, children)| async move { - let tree = ProtoTree { - root: Some(root_dir), - children: children.into(), - }; - let tree_digest = serialize_and_upload_message( - &tree, - cas_store.as_pin(), - &mut hasher.hasher(), - ) .await - .err_tip(|| format!("While processing {entry}"))?; - Ok(DirectoryInfo { - path: entry, - tree_digest, - }) - }) - .await - .err_tip(|| format!("Uploading directory {}", full_path.display()))?, - )) - } else if metadata.is_symlink() { - // Resolve the symlink to determine what it points to. - // Symlinks created by DirectoryCache (absolute paths into - // the cache directory) must NOT be uploaded as symlinks — - // the target path is worker-local and meaningless to the - // client. Instead, follow the symlink and upload the - // resolved content (file or directory). - let target = fs::read_link(&full_path).await.err_tip(|| { - format!("Reading symlink target for {}", full_path.display()) - })?; - let is_absolute_symlink = Path::new(&target).is_absolute(); - - if is_absolute_symlink { - // Absolute symlink — resolve and upload contents. - match fs::metadata(&full_path).await { - Ok(resolved_meta) => { - if resolved_meta.is_dir() { - // Upload as directory (Tree proto). - Ok(OutputType::Directory( - upload_directory( - cas_store.as_pin(), - &full_path, - work_directory, - hasher, - digest_uploaders, - ) - .and_then(|(root_dir, children)| async move { - let tree = ProtoTree { - root: Some(root_dir), - children: children.into(), - }; - let tree_digest = serialize_and_upload_message( - &tree, + .err_tip(|| format!("Uploading directory {}", full_path.display()))?, + )) + } else if metadata.is_symlink() { + // Resolve the symlink to determine what it points to. + // Symlinks created by DirectoryCache (absolute paths into + // the cache directory) must NOT be uploaded as symlinks — + // the target path is worker-local and meaningless to the + // client. Instead, follow the symlink and upload the + // resolved content (file or directory). + let target = fs::read_link(&full_path).await.err_tip(|| { + format!("Reading symlink target for {}", full_path.display()) + })?; + let is_absolute_symlink = Path::new(&target).is_absolute(); + + if is_absolute_symlink { + // Absolute symlink — resolve and upload contents. + match fs::metadata(&full_path).await { + Ok(resolved_meta) => { + if resolved_meta.is_dir() { + // Upload as directory (Tree proto). + Ok(OutputType::Directory( + upload_directory( cas_store.as_pin(), - &mut hasher.hasher(), + &full_path, + work_directory, + hasher, + digest_uploaders, + uploads, ) - .await - .err_tip(|| format!("While processing {entry}"))?; - Ok(DirectoryInfo { - path: entry, - tree_digest, + .and_then(|(root_dir, children)| async move { + let tree = ProtoTree { + root: Some(root_dir), + children: children.into(), + }; + let tree_digest = serialize_and_upload_message( + &tree, + cas_store.as_pin(), + &mut hasher.hasher(), + ) + .await + .err_tip(|| format!("While processing {entry}"))?; + Ok(DirectoryInfo { + path: entry, + tree_digest, + }) }) - }) - .await - .err_tip(|| { - format!( - "Uploading symlinked directory {}", - full_path.display() + .await + .err_tip( + || { + format!( + "Uploading symlinked directory {}", + full_path.display() + ) + }, + )?, + )) + } else { + // Upload as file (follow symlink). + Ok(OutputType::File( + upload_file( + cas_store.as_pin(), + &full_path, + hasher, + resolved_meta, + digest_uploaders, + uploads, ) - })?, - )) - } else { - // Upload as file (follow symlink). - Ok(OutputType::File( - upload_file( - cas_store.as_pin(), - &full_path, - hasher, - resolved_meta, - digest_uploaders, - ) - .await - .map(|mut file_info| { - file_info.name_or_path = NameOrPath::Path(entry); - file_info - }) - .err_tip(|| { + .await + .map(|mut file_info| { + file_info.name_or_path = NameOrPath::Path(entry); + file_info + }) + .err_tip( + || { + format!( + "Uploading symlinked file {}", + full_path.display() + ) + }, + )?, + )) + } + } + Err(e) => { + if e.code != Code::NotFound { + return Err(e).err_tip(|| { format!( - "Uploading symlinked file {}", + "While resolving absolute symlink {}", full_path.display() ) - })?, - )) + }); + } + Ok(OutputType::None) } } - Err(e) => { - if e.code != Code::NotFound { - return Err(e).err_tip(|| { - format!( - "While resolving absolute symlink {}", - full_path.display() - ) - }); + } else { + // Relative symlink — action intentionally created it. + // Upload as a proper symlink. + let output_symlink = upload_symlink(&full_path, work_directory) + .await + .map(|mut symlink_info| { + symlink_info.name_or_path = NameOrPath::Path(entry); + symlink_info + }) + .err_tip(|| format!("Uploading symlink {}", full_path.display()))?; + match fs::metadata(&full_path).await { + Ok(metadata) => { + if metadata.is_dir() { + Ok(OutputType::DirectorySymlink(output_symlink)) + } else { + // Note: If it's anything but directory we put it as a file symlink. + Ok(OutputType::FileSymlink(output_symlink)) + } } - Ok(OutputType::None) - } - } - } else { - // Relative symlink — action intentionally created it. - // Upload as a proper symlink. - let output_symlink = upload_symlink(&full_path, work_directory) - .await - .map(|mut symlink_info| { - symlink_info.name_or_path = NameOrPath::Path(entry); - symlink_info - }) - .err_tip(|| format!("Uploading symlink {}", full_path.display()))?; - match fs::metadata(&full_path).await { - Ok(metadata) => { - if metadata.is_dir() { - Ok(OutputType::DirectorySymlink(output_symlink)) - } else { - // Note: If it's anything but directory we put it as a file symlink. + Err(e) => { + if e.code != Code::NotFound { + return Err(e).err_tip(|| { + format!( + "While querying target symlink metadata for {}", + full_path.display() + ) + }); + } + // If the file doesn't exist, we consider it a file. Even though the + // file doesn't exist we still need to populate an entry. Ok(OutputType::FileSymlink(output_symlink)) } } - Err(e) => { - if e.code != Code::NotFound { - return Err(e).err_tip(|| { - format!( - "While querying target symlink metadata for {}", - full_path.display() - ) - }); - } - // If the file doesn't exist, we consider it a file. Even though the - // file doesn't exist we still need to populate an entry. - Ok(OutputType::FileSymlink(output_symlink)) - } } + } else { + Err(make_err!( + Code::Internal, + "{full_path:?} was not a file, folder or symlink. Must be one.", + )) } - } else { - Err(make_err!( - Code::Internal, - "{full_path:?} was not a file, folder or symlink. Must be one.", - )) } - }); + .boxed(), + ); } let mut output_files = vec![]; let mut output_folders = vec![]; @@ -2443,12 +3472,10 @@ impl RunningActionImpl { let mut output_file_symlinks = vec![]; if execution_result.exit_code != 0 { - let stdout = core::str::from_utf8(&execution_result.stdout).unwrap_or(""); - let stderr = core::str::from_utf8(&execution_result.stderr).unwrap_or(""); error!( exit_code = ?execution_result.exit_code, - stdout = ?stdout[..min(stdout.len(), 1000)], - stderr = ?stderr[..min(stderr.len(), 1000)], + stdout = ?log_excerpt(execution_result.stdout.head(), 1000), + stderr = ?log_excerpt(execution_result.stderr.head(), 1000), command = ?command_proto.arguments, "Command returned non-zero exit code", ); @@ -2458,11 +3485,15 @@ impl RunningActionImpl { let start = std::time::Instant::now(); let data = execution_result.stdout; let data_len = data.len(); - let digest = compute_buf_digest(&data, &mut hasher.hasher()); - cas_store - .update_oneshot(digest, data) - .await - .err_tip(|| "Uploading stdout")?; + // Boxed: the upload future is large, and two of them inline + // pushed the debug test binary past its thread stack. + let digest = Box::pin(upload_captured_output( + cas_store.as_pin(), + data, + hasher, + "stdout", + )) + .await?; debug!( ?digest, data_len, @@ -2475,11 +3506,15 @@ impl RunningActionImpl { let start = std::time::Instant::now(); let data = execution_result.stderr; let data_len = data.len(); - let digest = compute_buf_digest(&data, &mut hasher.hasher()); - cas_store - .update_oneshot(digest, data) - .await - .err_tip(|| "Uploading stderr")?; + // Boxed: the upload future is large, and two of them inline + // pushed the debug test binary past its thread stack. + let digest = Box::pin(upload_captured_output( + cas_store.as_pin(), + data, + hasher, + "stderr", + )) + .await?; debug!( ?digest, data_len, @@ -2494,6 +3529,8 @@ impl RunningActionImpl { num_output_paths = output_path_futures.len(), "upload_results: starting stdout/stderr/output_paths uploads", ); + let mut output_path_futures = + futures::stream::iter(output_path_futures).buffer_unordered(UPLOAD_CONCURRENCY); let join_start = std::time::Instant::now(); let upload_result = futures::try_join!(stdout_digest_fut, stderr_digest_fut, async { while let Some(output_type) = output_path_futures.try_next().await? { @@ -2619,12 +3656,61 @@ impl RunningAction for RunningActionImpl { } async fn prepare_action(self: Arc) -> Result, Error> { - let res = self - .metrics() - .clone() - .prepare_action - .wrap(Self::inner_prepare_action(self)) - .await; + let download_timeout = self.running_actions_manager.max_download_timeout; + let operation_id = self.operation_id.clone(); + let metrics = self.metrics().clone(); + let mut kill_rx = self.kill_token.subscribe(); + // Boxed: the prepare state machine is large, and pinning it on the + // stack next to the select below overflowed a test thread. + let mut prepare_fut = Box::pin( + metrics + .prepare_action + .wrap(Self::inner_prepare_action(self)), + ); + // A fetch that never answers would otherwise hold the slot for good + // while the worker's keepalives say it is fine: name it every minute + // and give up at the timeout, so the scheduler can retry elsewhere. + let stall_warn_fut = async { + let mut elapsed_secs = 0u64; + loop { + tokio::time::sleep(Duration::from_mins(1)).await; + elapsed_secs += 60; + warn!( + ?operation_id, + elapsed_s = elapsed_secs, + timeout_s = download_timeout.as_secs(), + "prepare_action: still fetching inputs, possible stall", + ); + } + }; + let res = tokio::time::timeout(download_timeout, async { + tokio::pin!(stall_warn_fut); + tokio::select! { + result = &mut prepare_fut => result, + _ = kill_rx.wait_for(|killed| *killed) => { + warn!(%operation_id, "prepare_action: killed while fetching inputs"); + Err(make_err!( + Code::Aborted, + "Operation {operation_id} was killed while fetching its inputs" + )) + } + () = &mut stall_warn_fut => unreachable!(), + } + }) + .await + .map_err(|_| { + warn!( + %operation_id, + timeout_s = download_timeout.as_secs(), + "prepare_action: input fetch timed out" + ); + make_err!( + Code::DeadlineExceeded, + "Fetching the inputs of operation {operation_id} took longer than max_download_timeout ({}s)", + download_timeout.as_secs(), + ) + }) + .and_then(|res| res); if let Err(ref e) = res { warn!(?e, "Error during prepare_action"); } @@ -2653,6 +3739,7 @@ impl RunningAction for RunningActionImpl { "upload_results: starting with timeout", ); let metrics = self.metrics().clone(); + let kill_token = self.kill_token.clone(); let upload_fut = metrics .upload_results .wrap(Self::inner_upload_results(self)); @@ -2671,11 +3758,24 @@ impl RunningAction for RunningActionImpl { } }; + // Only a kill that arrives during the upload ends it: a kill that + // already ended the command leaves a killed result that still has to + // reach the scheduler, so the token's current state is taken as seen + // and only a change from here on aborts. + let mut kill_rx = kill_token.subscribe(); + kill_rx.borrow_and_update(); let res = tokio::time::timeout(upload_timeout, async { tokio::pin!(upload_fut); tokio::pin!(stall_warn_fut); tokio::select! { result = &mut upload_fut => result, + Ok(()) = kill_rx.changed() => { + warn!(%operation_id, "upload_results: killed while uploading"); + Err(make_err!( + Code::Aborted, + "Operation {operation_id} was killed while uploading its results" + )) + } () = &mut stall_warn_fut => unreachable!(), } }) @@ -2793,8 +3893,98 @@ impl Debug for Callbacks { /// may be used to run the action with a particular set of additional /// environment variables, or perhaps configure it to execute within a /// container. +/// What the worker holds an action to beyond the scheduler's ledger. Each +/// axis is on or off on its own; the value exists when at least one is on. +#[derive(Debug, Clone)] +pub struct ResourceEnforcement { + /// Soft memory enforcement: the worker kills an action whose process + /// group grows past its reservation plus headroom. Linux only; the + /// sampler that measures it does not exist elsewhere. An action run + /// through a persistent worker is not covered: that process outlives + /// the action and serves others, so its memory is no one action's + /// reservation, and only the pod's limit bounds it. + pub memory: Option, + /// Refuse an action whose disk reservation, under this property, is + /// more than the free space under the work directory less what the + /// actions already admitted reserved, and measure what the action + /// writes. + pub disk_property_name: Option, + /// Kill an action whose own files pass the disk reservation plus + /// `disk_headroom_percent`. + pub disk_soft: bool, + pub disk_headroom_percent: u64, +} + +/// The reservation a memory ceiling is built from. +#[derive(Debug, Clone)] +pub struct MemoryLimit { + /// Platform property carrying the reservation in KiB. + pub property_name: String, + /// Percent above the reservation the action may reach. + pub headroom_percent: u64, +} + +impl ResourceEnforcement { + /// From the worker's configuration. `None` when neither axis is on; + /// `disk: guard` alone is honoured as much as `memory: soft` alone. + pub fn from_config(config: &ResourceEnforcementConfig) -> Option { + let memory = match config.memory { + MemoryEnforcement::None => None, + MemoryEnforcement::Soft => Some(MemoryLimit { + property_name: config.memory_property_name.clone(), + headroom_percent: config.memory_headroom_percent, + }), + }; + let disk_property_name = + (config.disk != DiskEnforcement::None).then(|| config.disk_property_name.clone()); + (memory.is_some() || disk_property_name.is_some()).then_some(Self { + memory, + disk_property_name, + disk_soft: config.disk == DiskEnforcement::Soft, + disk_headroom_percent: config.disk_headroom_percent, + }) + } +} + +/// Free space in KiB on the filesystem holding `path`, as an unprivileged +/// process may use it. +#[cfg(target_family = "unix")] +pub fn free_disk_kb(path: &str) -> Option { + let c_path = std::ffi::CString::new(path).ok()?; + let mut stat: libc::statvfs = unsafe { core::mem::zeroed() }; + // SAFETY: `c_path` is a valid C string and `stat` is a zeroed out + // parameter the call fills in. + let rc = unsafe { libc::statvfs(c_path.as_ptr(), &raw mut stat) }; + if rc != 0 { + return None; + } + Some(widen(stat.f_bavail).saturating_mul(widen(stat.f_frsize)) / 1024) +} + +/// `statvfs` fields are u32 or u64 depending on the platform. +#[cfg(target_family = "unix")] +fn widen(value: impl Into) -> u64 { + value.into() +} + +#[cfg(not(target_family = "unix"))] +pub fn free_disk_kb(_path: &str) -> Option { + None +} + #[derive(Debug, Default)] pub struct ExecutionConfiguration { + /// Bytes of stdout or stderr kept in memory before the rest spills to a + /// file under the action directory; 0 keeps everything in memory. + pub max_captured_output_bytes: u64, + /// Kill an action that exceeds its memory reservation, before the + /// pod's cgroup limit takes the whole worker with it. + pub resource_enforcement: Option, + /// Time a timed-out or cancelled action gets between SIGTERM and + /// SIGKILL to write its own cleanup; zero sends SIGKILL at once. + pub kill_grace: Duration, + /// Point `TMPDIR` at a directory of the action's own, removed with it. + pub set_tmpdir: bool, /// Buck2-only helper, configured for a dedicated execution container. pub buck2_file_capture: Option, /// If set, will be executed instead of the first argument passed in the @@ -2805,6 +3995,25 @@ pub struct ExecutionConfiguration { /// executes other than those in the `ActionInfo`. On Windows, `SystemRoot` /// and PATH are also assigned (see `inner_execute`). pub additional_environment: Option>, + /// Bazel persistent workers: whether an action that supports them runs + /// in a pooled process, and how the pool is sized. + pub persistent_workers: PersistentWorkersSettings, +} + +/// See `ExecutionConfiguration::persistent_workers`. +#[derive(Debug, Clone, Copy)] +pub struct PersistentWorkersSettings { + pub enabled: bool, + pub pool: PoolConfig, +} + +impl Default for PersistentWorkersSettings { + fn default() -> Self { + Self { + enabled: true, + pool: PoolConfig::default(), + } + } } #[derive(Debug)] @@ -3055,7 +4264,12 @@ impl UploadActionResults { pub enum UseNamespaces { No, Yes, - YesAndMount, + /// Also unshare the mount namespace. With `isolate_tmp` each action gets + /// a private `/tmp`, its own tmp directory bound over `/tmp`, on top of + /// the masked root action directory. + YesAndMount { + isolate_tmp: bool, + }, } #[derive(Debug)] @@ -3068,6 +4282,7 @@ pub struct RunningActionsManagerArgs<'a> { pub upload_action_result_config: &'a UploadActionResultConfig, pub max_action_timeout: Duration, pub max_upload_timeout: Duration, + pub max_download_timeout: Duration, pub max_cleanup_wait: Duration, pub max_cleanup_backoff: Duration, pub timeout_handled_externally: bool, @@ -3081,7 +4296,10 @@ pub struct RunningActionsManagerArgs<'a> { pub use_namespaces: UseNamespaces, } -struct CleanupGuard { +/// The mark that an operation's directory is being cleaned; dropping it +/// clears the mark and wakes whoever waits for it. +#[derive(Debug)] +pub struct CleanupGuard { manager: Weak, operation_id: OperationId, } @@ -3108,7 +4326,11 @@ pub struct RunningActionsManagerImpl { upload_action_results: UploadActionResults, max_action_timeout: Duration, max_upload_timeout: Duration, + max_download_timeout: Duration, timeout_handled_externally: bool, + /// The container's memory limit, the yardstick for calling a SIGKILL + /// nobody here sent an OOM kill. + pod_memory_limit_kb: Option, #[cfg(target_os = "linux")] use_namespaces: UseNamespaces, running_actions: Mutex>>, @@ -3117,6 +4339,10 @@ pub struct RunningActionsManagerImpl { action_done_tx: watch::Sender<()>, callbacks: Callbacks, metrics: Arc, + /// Disk each admitted action reserved, by operation, so the disk guard + /// checks against what is left after them. An entry lives from the + /// guard passing to `cleanup_action`. + disk_reservations_kb: Mutex>, /// Track operations being cleaned up to avoid directory collisions during action retries. /// When an action fails and is retried on the same worker, we need to ensure the previous /// attempt's directory is fully cleaned up before creating a new one. @@ -3134,6 +4360,8 @@ pub struct RunningActionsManagerImpl { /// CAS tiers (opt-in via `experimental_active_input_leases`). active_input_leases: bool, persistent_worker_pool: PersistentWorkerPool, + /// Retires idle persistent workers; ends with the manager. + _persistent_worker_sweeper: JoinHandleDropGuard<()>, buck2_captures: Mutex>, } @@ -3153,6 +4381,9 @@ impl RunningActionsManagerImpl { .get_arc() .err_tip(|| "FilesystemStore's internal Arc was lost")?; let (action_done_tx, _) = watch::channel(()); + let persistent_worker_pool = + PersistentWorkerPool::new(args.execution_configuration.persistent_workers.pool); + let persistent_worker_sweeper = persistent_worker_pool.spawn_sweeper(); Ok(Self { root_action_directory: args.root_action_directory, execution_configuration: args.execution_configuration, @@ -3165,7 +4396,9 @@ impl RunningActionsManagerImpl { ) .err_tip(|| "During RunningActionsManagerImpl construction")?, max_action_timeout: args.max_action_timeout, + pod_memory_limit_kb: pod_memory_limit_kb(), max_upload_timeout: args.max_upload_timeout, + max_download_timeout: args.max_download_timeout, timeout_handled_externally: args.timeout_handled_externally, running_actions: Mutex::new(HashMap::new()), action_done_tx, @@ -3174,13 +4407,15 @@ impl RunningActionsManagerImpl { directory_cache: args.directory_cache.as_ref().map(Arc::downgrade), ..Default::default() }), + disk_reservations_kb: Mutex::new(HashMap::new()), cleaning_up_operations: Mutex::new(HashSet::new()), max_cleanup_wait: args.max_cleanup_wait, max_cleanup_backoff: args.max_cleanup_backoff, cleanup_complete_notify: Arc::new(Notify::new()), directory_cache: args.directory_cache, active_input_leases: args.active_input_leases, - persistent_worker_pool: PersistentWorkerPool::default(), + persistent_worker_pool, + _persistent_worker_sweeper: persistent_worker_sweeper, buck2_captures: Mutex::new(HashMap::new()), #[cfg(target_os = "linux")] use_namespaces: args.use_namespaces, @@ -3200,18 +4435,20 @@ impl RunningActionsManagerImpl { /// Fixes a race condition that occurs when an action fails to execute on a worker, and the same worker /// attempts to re-execute the same action before the physical cleanup (file is removed) completes. /// See this issue for additional details: - async fn wait_for_cleanup_if_needed(&self, operation_id: &OperationId) -> Result<(), Error> { + async fn wait_for_cleanup_if_needed( + self: &Arc, + operation_id: &OperationId, + ) -> Result<(), Error> { let start = Instant::now(); let mut backoff = Duration::from_millis(10); let mut has_waited = false; loop { - let should_wait = { - let cleaning = self.cleaning_up_operations.lock(); - cleaning.contains(operation_id) - }; - - if !should_wait { + // The mark is held for the removal below, so it cannot race the + // orphan sweep's removal of the same tree; while someone else + // holds it, this waits. + let mark = self.perform_cleanup(operation_id.clone()); + if let Some(_cleaning_guard) = mark { let action_is_running = self .running_actions .lock() @@ -3342,18 +4579,118 @@ impl RunningActionsManagerImpl { let action = get_and_decode_digest::(self.cas_store.as_ref(), action_digest.into()) .await - .err_tip(|| "During start_action")?; - let action_info = ActionInfo::try_from_action_and_execute_request( + .err_tip(|| "During start_action") + .map_err(tag_missing_input)?; + let mut action_info = ActionInfo::try_from_action_and_execute_request( execute_request, action, load_start_timestamp, queued_timestamp, ) .err_tip(|| "Could not create ActionInfo in create_and_add_action()")?; + // With enforcement on, the scheduler's view of the properties + // replaces the client's: a hint or cold-start reservation the + // scheduler placed this action by only exists there, and what + // is enforced (and exported to the environment) has to be the + // number the ledger charged. Off, the worker keeps reading the + // client's copy as it always did, so an upgrade changes no + // action's environment on its own. + if self.execution_configuration.resource_enforcement.is_some() + && let Some(platform) = start_execute.platform + { + for property in platform.properties { + action_info + .platform_properties + .insert(property.name, property.value); + } + } Ok(action_info) }) } + /// Removes action directories under the root that no running action + /// owns and that have not been touched for `max_cleanup_wait`. A + /// cleanup that failed (the store was down, the process restarted + /// mid-removal) leaves its directory behind, and it counts against the + /// pod's disk until something removes it. A directory that will not go + /// is logged and counted and the sweep carries on, or one stubborn + /// entry would stop every sweep at the same place. Returns how many + /// were removed; the error is for the listing itself failing. + /// + /// A directory is named for its operation, and a retry of that + /// operation may be starting at any moment. The retry's + /// `wait_for_cleanup_if_needed` waits while the operation is marked as + /// cleaning, so the sweep takes that mark for the removal and reads + /// ownership and age under it; a mark already held is the action's own + /// cleanup, which finishes the job. + pub async fn sweep_orphaned_directories(self: &Arc) -> Result { + let (_permit, dir_handle) = fs::read_dir(&self.root_action_directory) + .await + .err_tip(|| format!("Reading {} for orphans", self.root_action_directory))? + .into_inner(); + let mut dir_stream = ReadDirStream::new(dir_handle); + let now = SystemTime::now(); + let mut removed = 0; + while let Some(entry) = dir_stream.next().await { + let entry = entry.err_tip(|| "Iterating action root for orphans")?; + let Some(name) = entry.file_name().to_str().map(str::to_owned) else { + continue; + }; + // A running action's directory is not the sweep's to touch, and + // its cleanup waits on the mark, so the mark is taken only for + // an unowned directory and given back at once if the directory + // became owned in between: an action that starts and ends + // inside that window then cleans up as usual, not later. + let is_owned = || { + self.running_actions + .lock() + .keys() + .any(|operation_id| operation_id.to_string() == name) + }; + if is_owned() { + continue; + } + let Some(_cleaning_guard) = self.perform_cleanup(OperationId::from(name.as_str())) + else { + continue; + }; + if is_owned() { + continue; + } + let path = entry.path(); + // Not followed: a symlink under the root is not an action + // directory, whatever it points at. + let Ok(metadata) = fs::symlink_metadata(&path).await else { + continue; + }; + if !metadata.is_dir() { + continue; + } + // Young directories may belong to an action between creation and + // registration; only settled ones are orphans. + let age = metadata + .modified() + .ok() + .and_then(|modified| now.duration_since(modified).ok()) + .unwrap_or(Duration::MAX); + if age < self.max_cleanup_wait { + continue; + } + warn!(path = %path.display(), age_secs = age.as_secs(), "Removing orphaned action directory"); + match fs::remove_dir_all(&path).await { + Ok(()) => { + self.metrics.orphan_removals.inc(); + removed += 1; + } + Err(err) => { + self.metrics.orphan_removal_failures.inc(); + warn!(path = %path.display(), ?err, "Could not remove orphaned action directory"); + } + } + } + Ok(removed) + } + fn cleanup_action(&self, operation_id: &OperationId) -> Result<(), Error> { // The guard must be dropped before notifying: `send_modify` takes // the watch channel's internal lock, and a `wait_for` closure that @@ -3369,6 +4706,7 @@ impl RunningActionsManagerImpl { ) }) }; + self.disk_reservations_kb.lock().remove(operation_id); // No need to copy anything, we just are telling the receivers an event happened. self.action_done_tx.send_modify(|()| {}); result.map(|_| ()) @@ -3381,6 +4719,9 @@ impl RunningActionsManagerImpl { operation_id = ?action.operation_id, "Sending kill to running operation", ); + // The token first: a phase waiting on it (a fetch, an upload) ends + // now, whether or not the command ever started. + action.kill_token.send_replace(true); let kill_channel_tx = { let mut action_state = action.state.lock(); action_state.kill_channel_tx.take() @@ -3395,14 +4736,41 @@ impl RunningActionsManagerImpl { } } - fn perform_cleanup(self: &Arc, operation_id: OperationId) -> Option { - let mut cleaning = self.cleaning_up_operations.lock(); - cleaning - .insert(operation_id.clone()) - .then_some(CleanupGuard { - manager: Arc::downgrade(self), - operation_id, - }) + /// Marks the operation as being cleaned and returns the guard that + /// clears the mark, or `None` when another cleanup already holds it. + /// + /// The guard is built only after the lock is released: its `Drop` takes + /// the same lock, so a guard made under it and thrown away (as + /// `then_some` did, whatever the insert said) deadlocked the thread on + /// its own mutex and, one by one, every runtime thread that then came to + /// clean up or start an action, while the worker's keepalives kept it + /// looking alive. The orphan sweep made the double mark routine: it marks + /// every directory it visits, and a directory whose action is mid-cleanup + /// is already marked. + pub fn perform_cleanup(self: &Arc, operation_id: OperationId) -> Option { + let inserted = self + .cleaning_up_operations + .lock() + .insert(operation_id.clone()); + inserted.then(|| CleanupGuard { + manager: Arc::downgrade(self), + operation_id, + }) + } + + /// Takes the cleanup mark, waiting for whoever holds it to let go. The + /// wait is armed before the mark is tried, so a release in between is + /// not missed. + pub async fn take_cleanup_mark(self: &Arc, operation_id: OperationId) -> CleanupGuard { + loop { + let notified = self.cleanup_complete_notify.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + if let Some(guard) = self.perform_cleanup(operation_id.clone()) { + return guard; + } + notified.await; + } } } @@ -3453,14 +4821,19 @@ impl RunningActionsManager for RunningActionsManagerImpl { } else { action_info.timeout }; - if timeout > self.max_action_timeout { - return Err(make_err!( - Code::InvalidArgument, - "Action timeout of {} seconds is greater than the maximum allowed timeout of {} seconds", - timeout.as_secs_f32(), - self.max_action_timeout.as_secs_f32() - )); - } + // REAPI lets a server clamp a timeout over its maximum; the + // client asked for more time, not for a rejection. + let timeout = if timeout > self.max_action_timeout { + info!( + ?operation_id, + requested_s = timeout.as_secs_f32(), + max_s = self.max_action_timeout.as_secs_f32(), + "Action timeout clamped to the worker maximum" + ); + self.max_action_timeout + } else { + timeout + }; let running_action = Arc::new(RunningActionImpl::new( execution_metadata, operation_id.clone(), @@ -3607,6 +4980,10 @@ pub struct Metrics { cleanup_waits: CounterWithTime, #[metric(help = "Number of stale directories removed during action retries.")] stale_removals: CounterWithTime, + #[metric(help = "Orphaned action directories removed by the periodic sweep.")] + orphan_removals: CounterWithTime, + #[metric(help = "Orphaned action directories the sweep could not remove.")] + orphan_removal_failures: CounterWithTime, #[metric(help = "Number of timeouts while waiting for cleanup to complete.")] cleanup_wait_timeouts: CounterWithTime, #[metric(help = "Stats about the get_proto_command_from_store command.")] @@ -3629,6 +5006,18 @@ pub struct Metrics { upload_stderr: AsyncCounterWrapper, #[metric(help = "Total number of task timeouts.")] task_timeouts: CounterWithTime, + #[metric(help = "Actions killed for exceeding their memory reservation.")] + memory_reservation_kills: CounterWithTime, + #[metric(help = "Actions refused because their disk reservation exceeded the free space.")] + disk_guard_refusals: CounterWithTime, + #[metric(help = "Actions killed for writing past their disk reservation.")] + disk_reservation_kills: CounterWithTime, + #[metric(help = "Actions served by a pooled persistent worker process.")] + persistent_worker_dispatches: CounterWithTime, + #[metric( + help = "Actions that asked for a persistent worker and ran one-shot instead: pool at capacity, spawn failure or an unsupported protocol." + )] + persistent_worker_fallbacks: CounterWithTime, #[metric( help = "Stats about the input-directory cache (hits, misses, subtree reuse, evictions, size)." )] diff --git a/nativelink-worker/src/worker_api_client_wrapper.rs b/nativelink-worker/src/worker_api_client_wrapper.rs index 1e2791fc093..cbe04df14eb 100644 --- a/nativelink-worker/src/worker_api_client_wrapper.rs +++ b/nativelink-worker/src/worker_api_client_wrapper.rs @@ -19,7 +19,7 @@ use nativelink_error::{make_err, Error, ResultExt}; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::update_for_scheduler::Update; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::worker_api_client::WorkerApiClient; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{ - ConnectWorkerRequest, ExecuteComplete, ExecuteResult, GoingAwayRequest, KeepAliveRequest, UpdateForScheduler, UpdateForWorker + ConnectWorkerRequest, ExecuteAccepted, ExecuteComplete, ExecuteDeclined, ExecuteResult, GoingAwayRequest, KeepAliveRequest, UpdateForScheduler, UpdateForWorker }; use tokio::sync::mpsc::Sender; use tonic::codec::Streaming; @@ -49,6 +49,16 @@ pub trait WorkerApiClientTrait: Clone + Sync + Send + Sized + Unpin { request: ExecuteResult, ) -> impl Future> + Send; + fn execute_accepted( + &mut self, + request: ExecuteAccepted, + ) -> impl Future> + Send; + + fn execute_declined( + &mut self, + request: ExecuteDeclined, + ) -> impl Future> + Send; + fn execution_complete( &mut self, request: ExecuteComplete, @@ -133,4 +143,12 @@ impl WorkerApiClientTrait for WorkerApiClientWrapper { async fn execution_complete(&mut self, request: ExecuteComplete) -> Result<(), Error> { self.send_update(Update::ExecuteComplete(request)).await } + + async fn execute_accepted(&mut self, request: ExecuteAccepted) -> Result<(), Error> { + self.send_update(Update::ExecuteAccepted(request)).await + } + + async fn execute_declined(&mut self, request: ExecuteDeclined) -> Result<(), Error> { + self.send_update(Update::ExecuteDeclined(request)).await + } } diff --git a/nativelink-worker/tests/capacity_test.rs b/nativelink-worker/tests/capacity_test.rs index bf84f8426ed..a3a0b8cdeb0 100644 --- a/nativelink-worker/tests/capacity_test.rs +++ b/nativelink-worker/tests/capacity_test.rs @@ -1,8 +1,194 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; +use std::path::{Path, PathBuf}; + +use nativelink_config::cas_server::{ + CapacityConfig, CpuUnit, DiskEnforcement, MemoryEnforcement, ResourceEnforcementConfig, + WorkerProperty, +}; +use nativelink_error::Code; +#[cfg(target_family = "unix")] +use nativelink_worker::capacity::find_limit; use nativelink_worker::capacity::{ - free_memory_kb_from, parse_meminfo_available_kb, parse_memory_current_kb, + ObservedCapacity, advertised, free_memory_kb_from, memory_headroom_percent, parse_cpu_max, + parse_meminfo_available_kb, parse_meminfo_total_kb, parse_memory_current_kb, parse_memory_max, + parse_self_cgroup, resolve_cgroup_dir, without_cgroup, }; use pretty_assertions::assert_eq; +fn config() -> CapacityConfig { + config_in(CpuUnit::Cores) +} + +fn config_in(cpu_unit: CpuUnit) -> CapacityConfig { + CapacityConfig { + cpu_property_name: "cpu_count".to_string(), + memory_property_name: "memory_kb".to_string(), + cpu_unit, + overhead_cpu_millicores: 1000, + overhead_memory_kb: 4 * 1024 * 1024, + memory_headroom_percent: None, + } +} + +fn enforcement(memory: MemoryEnforcement, headroom: u64) -> ResourceEnforcementConfig { + ResourceEnforcementConfig { + memory, + disk: DiskEnforcement::None, + disk_property_name: "disk_kb".to_string(), + memory_property_name: "memory_kb".to_string(), + memory_headroom_percent: headroom, + disk_headroom_percent: 20, + } +} + +#[test] +fn cpu_max_is_quota_over_period_in_millicores() { + assert_eq!(parse_cpu_max("1400000 100000\n", 64), Some(14_000)); + assert_eq!(parse_cpu_max("max 100000\n", 8), Some(8_000)); + assert_eq!(parse_cpu_max("garbage", 8), None); + assert_eq!(parse_cpu_max("100000 0", 8), None); +} + +#[test] +fn memory_max_is_bytes_or_the_host() { + assert_eq!(parse_memory_max("55834574848\n", 1), Some(54_525_952)); + assert_eq!(parse_memory_max("max\n", 65_536_000), Some(65_536_000)); + assert_eq!(parse_memory_max("", 1), None); + assert_eq!( + parse_meminfo_total_kb("MemTotal: 65536000 kB\nMemFree: 1 kB\n"), + Some(65_536_000) + ); +} + +/// The drydock pool: a 52 GiB, 14 core pod, 4 GiB and one core kept back, +/// 20% enforcement headroom, comes out at 13 cores and 40 GiB, which is what +/// the Terraform sizing derives by hand today. +#[test] +fn advertised_takes_off_the_overhead_and_the_headroom() { + let observed = ObservedCapacity { + cpu_millicores: 14_000, + memory_kb: 52 * 1024 * 1024, + }; + assert_eq!(advertised(observed, &config(), 20), (13, 40 * 1024 * 1024)); + assert_eq!(advertised(observed, &config(), 0), (13, 48 * 1024 * 1024)); +} + +/// The headroom is typed once: unset, the block follows the enforcement +/// headroom while memory enforcement is on, and is nothing otherwise; set, +/// its own figure wins. +#[test] +fn headroom_follows_enforcement_unless_set() { + let unset = config(); + assert_eq!( + memory_headroom_percent(&unset, Some(&enforcement(MemoryEnforcement::Soft, 20))), + 20 + ); + assert_eq!( + memory_headroom_percent(&unset, Some(&enforcement(MemoryEnforcement::None, 20))), + 0 + ); + assert_eq!(memory_headroom_percent(&unset, None), 0); + let set = CapacityConfig { + memory_headroom_percent: Some(10), + ..config() + }; + assert_eq!( + memory_headroom_percent(&set, Some(&enforcement(MemoryEnforcement::Soft, 20))), + 10 + ); +} + +/// Without a readable cgroup the configured properties stand if they carry +/// both numbers; a worker that would register with neither is refused, +/// naming what is missing, rather than idling forever on a warning. +#[test] +fn without_a_cgroup_the_worker_keeps_its_properties_or_refuses_to_start() { + let both: HashMap = HashMap::from([ + ( + "cpu_count".to_string(), + WorkerProperty::Values(vec!["16".to_string()]), + ), + ( + "memory_kb".to_string(), + WorkerProperty::Values(vec!["60000000".to_string()]), + ), + ]); + assert!(without_cgroup(&config(), &both).is_ok()); + + let only_cpu: HashMap = HashMap::from([( + "cpu_count".to_string(), + WorkerProperty::Values(vec!["16".to_string()]), + )]); + let err = without_cgroup(&config(), &only_cpu).expect_err("memory_kb is missing"); + assert_eq!(err.code, Code::FailedPrecondition); + assert!(err.to_string().contains("no memory_kb"), "{err}"); + + let neither: HashMap = HashMap::new(); + let err = without_cgroup(&config(), &neither).expect_err("both are missing"); + assert!( + err.to_string().contains("no cpu_count or memory_kb"), + "{err}" + ); +} + +/// `cpu_count` is whole cores everywhere else in the configuration (`nproc`, +/// `cpu_count=1` per action), so that is the default scale: a 14-core pod +/// with the default block advertises 13, not the 13000 that would let the +/// scheduler pack thousands of actions onto it. A partial core rounds down. +/// A fleet on thousandths of a core says so and gets the exact figure. +#[test] +fn cpu_is_advertised_in_whole_cores_unless_told_millicores() { + let fourteen_cores = ObservedCapacity { + cpu_millicores: 14_000, + memory_kb: 1, + }; + let default_block = CapacityConfig { + cpu_property_name: "cpu_count".to_string(), + memory_property_name: "memory_kb".to_string(), + cpu_unit: CpuUnit::default(), + overhead_cpu_millicores: 1000, + overhead_memory_kb: 0, + memory_headroom_percent: None, + }; + assert_eq!(advertised(fourteen_cores, &default_block, 0).0, 13); + assert_eq!( + advertised(fourteen_cores, &config_in(CpuUnit::Millicores), 0).0, + 13_000 + ); + let fourteen_and_a_half = ObservedCapacity { + cpu_millicores: 14_500, + memory_kb: 1, + }; + assert_eq!(advertised(fourteen_and_a_half, &config(), 0).0, 13); + assert_eq!( + advertised(fourteen_and_a_half, &config_in(CpuUnit::Millicores), 0).0, + 13_500 + ); +} + +#[test] +fn overhead_larger_than_the_limit_advertises_zero_not_a_wraparound() { + let observed = ObservedCapacity { + cpu_millicores: 500, + memory_kb: 1024, + }; + assert_eq!(advertised(observed, &config(), 20), (0, 0)); +} + /// The keepalive reports the limit less current usage when there is a /// limit, and what the host has available when there is not. #[test] @@ -24,3 +210,110 @@ fn free_memory_is_limit_less_current_or_the_host_available() { assert_eq!(free_memory_kb_from(None, Some(200), Some(777)), Some(777)); assert_eq!(free_memory_kb_from(None, None, None), None); } + +#[test] +fn self_cgroup_is_the_v2_line() { + assert_eq!(parse_self_cgroup("0::/\n"), Some("/")); + assert_eq!( + parse_self_cgroup( + "0::/kubepods.slice/kubepods-burstable.slice/kubepods-burstable-pod5ed6.slice/cri-containerd-e0fd.scope\n" + ), + Some( + "/kubepods.slice/kubepods-burstable.slice/kubepods-burstable-pod5ed6.slice/cri-containerd-e0fd.scope" + ) + ); + // cgroup v1 lines carry a controller name and no `0::` line. + assert_eq!( + parse_self_cgroup("12:memory:/docker/abc\n1:name=systemd:/docker/abc\n"), + None + ); + assert_eq!(parse_self_cgroup(""), None); + assert_eq!(parse_self_cgroup("0::relative\n"), None); +} + +#[test] +fn cgroup_dir_is_the_own_path_under_the_mount_when_it_exists() { + let root = Path::new("/sys/fs/cgroup"); + let scope = "/kubepods.slice/kubepods-burstable-pod5ed6.slice/cri-containerd-e0fd.scope"; + let own = PathBuf::from( + "/sys/fs/cgroup/kubepods.slice/kubepods-burstable-pod5ed6.slice/cri-containerd-e0fd.scope", + ); + // A privileged container sees the host's tree: its own scope is a + // directory below the root. + assert_eq!( + resolve_cgroup_dir(root, Some(&format!("0::{scope}\n")), |dir| dir == own), + own + ); + // A container with its own cgroup namespace is at the root. + assert_eq!( + resolve_cgroup_dir(root, Some("0::/\n"), |dir| dir == root), + root + ); + // The path is not under the mount (a cgroup namespace whose root is + // deeper than the mount shows): the root stands. + assert_eq!( + resolve_cgroup_dir(root, Some(&format!("0::{scope}\n")), |_| false), + root + ); + // Nothing readable at all: the root, as before. + assert_eq!(resolve_cgroup_dir(root, None, |_| true), root); +} + +// The walk is Linux-only code, and the test's paths are unix strings. +#[cfg(target_family = "unix")] +#[test] +fn limit_comes_from_the_nearest_limited_ancestor() { + let root = Path::new("/sys/fs/cgroup"); + let own = root.join("kubepods.slice/kubepods-pod1.slice/cri-containerd-abc.scope"); + // The scope itself is unlimited; the pod slice above it carries the limit. + let read = |path: &Path| -> Option { + match path.to_str()? { + "/sys/fs/cgroup/kubepods.slice/kubepods-pod1.slice/cri-containerd-abc.scope/memory.max" => { + Some("max\n".to_string()) + } + "/sys/fs/cgroup/kubepods.slice/kubepods-pod1.slice/memory.max" => { + Some("55834574848\n".to_string()) + } + "/sys/fs/cgroup/kubepods.slice/kubepods-pod1.slice/cri-containerd-abc.scope/cpu.max" => { + Some("max 100000\n".to_string()) + } + "/sys/fs/cgroup/kubepods.slice/cpu.max" => Some("1400000 100000\n".to_string()), + _ => None, + } + }; + assert_eq!( + find_limit(&own, root, "memory.max", read), + Some(( + root.join("kubepods.slice/kubepods-pod1.slice"), + "55834574848\n".to_string() + )) + ); + // CPU and memory may be limited at different levels. + assert_eq!( + find_limit(&own, root, "cpu.max", read), + Some((root.join("kubepods.slice"), "1400000 100000\n".to_string())) + ); +} + +// The walk is Linux-only code, and the test's paths are unix strings. +#[cfg(target_family = "unix")] +#[test] +fn no_limit_at_any_level_is_no_limit() { + let root = Path::new("/sys/fs/cgroup"); + let own = root.join("system.slice/nativelink.service"); + // A bare-metal systemd host: every level says max. + let all_max = |_: &Path| Some("max\n".to_string()); + assert_eq!(find_limit(&own, root, "memory.max", all_max), None); + // Nothing readable at all. + assert_eq!(find_limit(&own, root, "memory.max", |_| None), None); + // The walk stops at the mount root, never above it. + let outside = |path: &Path| (path == Path::new("/sys/memory.max")).then(|| "1\n".to_string()); + assert_eq!(find_limit(&own, root, "memory.max", outside), None); + // A directory not under the mount root is only looked at itself. + let own_only = + |path: &Path| (path == Path::new("/elsewhere/memory.max")).then(|| "2048\n".to_string()); + assert_eq!( + find_limit(Path::new("/elsewhere"), root, "memory.max", own_only), + Some((PathBuf::from("/elsewhere"), "2048\n".to_string())) + ); +} diff --git a/nativelink-worker/tests/local_worker_test.rs b/nativelink-worker/tests/local_worker_test.rs index 22cf78aaa87..f9f4dc23fc1 100644 --- a/nativelink-worker/tests/local_worker_test.rs +++ b/nativelink-worker/tests/local_worker_test.rs @@ -41,7 +41,7 @@ use nativelink_proto::build::bazel::remote::execution::v2::platform::Property; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::update_for_worker::Update; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{ ConnectWorkerRequest, ConnectionResult, ExecuteResult, KillOperationRequest, StartExecute, - UpdateForWorker, execute_result, + UpdateForWorker, execute_declined, execute_result, }; use nativelink_store::fast_slow_store::FastSlowStore; use nativelink_store::filesystem_store::FilesystemStore; @@ -52,10 +52,13 @@ use nativelink_util::action_messages::{ }; use nativelink_util::common::{DigestInfo, encode_stream_proto, fs, make_temp_path}; use nativelink_util::digest_hasher::DigestHasherFunc; +use nativelink_util::health_utils::{HealthStatus, HealthStatusIndicator}; use nativelink_util::store_trait::Store; +use nativelink_worker::capacity::free_memory_kb; use nativelink_worker::local_worker::new_local_worker; #[cfg(target_family = "unix")] use nativelink_worker::local_worker::preconditions_met; +use nativelink_worker::running_actions_manager::MISSING_INPUT_ERROR_TIP; use pretty_assertions::assert_eq; use prost::Message; use tokio::io::AsyncWriteExt; @@ -176,6 +179,8 @@ async fn kill_all_called_on_disconnect() -> Result<(), Error> { encode_stream_proto(&UpdateForWorker { update: Some(Update::ConnectionResult(ConnectionResult { worker_id: "foobar".to_string(), + dispatch_ack: false, + memory_property: String::new(), })), }) .unwrap(), @@ -214,6 +219,8 @@ async fn reconnects_when_action_stuck_in_transit_on_disconnect() -> Result<(), E encode_stream_proto(&UpdateForWorker { update: Some(Update::ConnectionResult(ConnectionResult { worker_id: expected_worker_id.clone(), + dispatch_ack: false, + memory_property: String::new(), })), }) .unwrap(), @@ -297,6 +304,8 @@ async fn blake3_digest_function_registered_properly() -> Result<(), Error> { encode_stream_proto(&UpdateForWorker { update: Some(Update::ConnectionResult(ConnectionResult { worker_id: expected_worker_id.clone(), + dispatch_ack: false, + memory_property: String::new(), })), }) .unwrap(), @@ -413,6 +422,8 @@ async fn start_action_lifecycle_test(single_use: bool) -> Result<(), Error> { encode_stream_proto(&UpdateForWorker { update: Some(Update::ConnectionResult(ConnectionResult { worker_id: expected_worker_id.clone(), + dispatch_ack: false, + memory_property: String::new(), })), }) .unwrap(), @@ -741,6 +752,8 @@ async fn experimental_precondition_script_fails() -> Result<(), Error> { encode_stream_proto(&UpdateForWorker { update: Some(Update::ConnectionResult(ConnectionResult { worker_id: expected_worker_id.clone(), + dispatch_ack: false, + memory_property: String::new(), })), }) .unwrap(), @@ -830,6 +843,8 @@ async fn kill_action_request_kills_action() -> Result<(), Error> { encode_stream_proto(&UpdateForWorker { update: Some(Update::ConnectionResult(ConnectionResult { worker_id: expected_worker_id.clone(), + dispatch_ack: false, + memory_property: String::new(), })), }) .unwrap(), @@ -928,6 +943,8 @@ async fn cas_not_found_returns_failed_precondition_test() -> Result<(), Error> { encode_stream_proto(&UpdateForWorker { update: Some(Update::ConnectionResult(ConnectionResult { worker_id: expected_worker_id.clone(), + dispatch_ack: false, + memory_property: String::new(), })), }) .unwrap(), @@ -980,14 +997,13 @@ async fn cas_not_found_returns_failed_precondition_test() -> Result<(), Error> { .expect_create_and_add_action(Ok(running_action.clone())) .await; - // Simulate prepare_action failing with a CAS NotFound error containing the - // specific "not found in either fast or slow store" message. This is the exact - // condition that the code checks to decide whether to return FailedPrecondition. + // Simulate prepare_action failing the way the manager fails a missing + // input: NotFound tagged with the missing-input tip, whatever the store + // itself said. + let missing_input = make_err!(Code::NotFound, "Hash 0123456789abcdef not found") + .append(MISSING_INPUT_ERROR_TIP); running_action - .expect_prepare_action(Err(make_err!( - Code::NotFound, - "Hash 0123456789abcdef not found in either fast or slow store" - ))) + .expect_prepare_action(Err(missing_input.clone())) .await?; // Cleanup is still called even when prepare_action fails. @@ -1000,7 +1016,8 @@ async fn cas_not_found_returns_failed_precondition_test() -> Result<(), Error> { let expected_action_result = ActionResult { error: Some(make_err!( Code::FailedPrecondition, - "Hash 0123456789abcdef not found in either fast or slow store" + "{}", + missing_input.message_string() )), ..ActionResult::default() }; @@ -1041,6 +1058,8 @@ async fn non_cas_not_found_returns_internal_error_test() -> Result<(), Error> { encode_stream_proto(&UpdateForWorker { update: Some(Update::ConnectionResult(ConnectionResult { worker_id: expected_worker_id.clone(), + dispatch_ack: false, + memory_property: String::new(), })), }) .unwrap(), @@ -1130,7 +1149,12 @@ async fn preconditions_met_extra_envs() -> Result<(), Error> { // So we have bash for nix cases, because the PATH gets reset extra_envs.insert("PATH".into(), env::var("PATH").unwrap()); - preconditions_met(Some("bash -c \"echo $DEMO_ENV\"".to_string()), &extra_envs).await?; + preconditions_met( + Some("bash -c \"echo $DEMO_ENV\"".to_string()), + &extra_envs, + Duration::from_secs(30), + ) + .await?; assert!(logs_contain("test_value_for_demo_env")); Ok(()) } @@ -1163,6 +1187,8 @@ async fn keep_alive_fail_logs() -> Result<(), Error> { encode_stream_proto(&UpdateForWorker { update: Some(Update::ConnectionResult(ConnectionResult { worker_id: "foobar".to_string(), + dispatch_ack: false, + memory_property: String::new(), })), }) .unwrap(), @@ -1241,6 +1267,8 @@ async fn disconnect_with_action_in_transit(single_use: bool) -> Result<(), Error encode_stream_proto(&UpdateForWorker { update: Some(Update::ConnectionResult(ConnectionResult { worker_id: expected_worker_id.clone(), + dispatch_ack: false, + memory_property: String::new(), })), }) .unwrap(), @@ -1322,3 +1350,277 @@ async fn disconnect_with_action_in_transit(single_use: bool) -> Result<(), Error Ok(()) } + +/// With a scheduler that speaks the acknowledgement, a worker at its +/// `max_inflight_tasks` declines the next dispatch instead of running it, +/// and says why; the first dispatch was acknowledged before it ran. +#[nativelink_test] +async fn a_worker_at_capacity_declines_the_next_dispatch() -> Result<(), Error> { + let mut test_context = setup_local_worker_with_config(LocalWorkerConfig { + max_inflight_tasks: 1, + worker_api_endpoint: EndpointConfig { + timeout: Some(10000.), + ..Default::default() + }, + ..Default::default() + }) + .await; + let streaming_response = test_context.maybe_streaming_response.take().unwrap(); + test_context + .client + .expect_connect_worker(Ok(streaming_response)) + .await; + + let worker_id = "foobar".to_string(); + let tx_stream = test_context.maybe_tx_stream.take().unwrap(); + tx_stream + .send(Frame::data( + encode_stream_proto(&UpdateForWorker { + update: Some(Update::ConnectionResult(ConnectionResult { + worker_id: worker_id.clone(), + dispatch_ack: true, + memory_property: String::new(), + })), + }) + .unwrap(), + )) + .await + .map_err(|e| make_input_err!("Could not send : {:?}", e))?; + + let action_info = ActionInfo { + command_digest: DigestInfo::new([1u8; 32], 10), + input_root_digest: DigestInfo::new([2u8; 32], 10), + timeout: Duration::from_secs(1), + platform_properties: HashMap::new(), + priority: 0, + load_timestamp: SystemTime::UNIX_EPOCH, + insert_timestamp: SystemTime::UNIX_EPOCH, + unique_qualifier: ActionUniqueQualifier::Uncacheable(ActionUniqueKey { + execution_scope: None, + instance_name: INSTANCE_NAME.to_string(), + digest_function: DigestHasherFunc::Blake3, + digest: DigestInfo::new([3u8; 32], 10), + }), + }; + let start = |operation_id: &str| UpdateForWorker { + update: Some(Update::StartAction(StartExecute { + request_metadata: None, + execute_request: Some((&action_info).into()), + operation_id: operation_id.to_string(), + queued_timestamp: None, + platform: Some(Platform::default()), + worker_id: worker_id.clone(), + })), + }; + + // The first dispatch is acknowledged, then runs (and stays running: + // nothing answers its prepare). + tx_stream + .send(Frame::data(encode_stream_proto(&start("first")).unwrap())) + .await + .map_err(|e| make_input_err!("Could not send : {:?}", e))?; + let accepted = test_context.client.expect_execute_accepted(Ok(())).await; + assert_eq!(accepted.operation_id, "first"); + let running_action = Arc::new(MockRunningAction::new()); + test_context + .actions_manager + .expect_create_and_add_action(Ok(running_action.clone())) + .await; + + // The second finds the worker full and is declined with the reason. + tx_stream + .send(Frame::data(encode_stream_proto(&start("second")).unwrap())) + .await + .map_err(|e| make_input_err!("Could not send : {:?}", e))?; + let declined = test_context.client.expect_execute_declined(Ok(())).await; + assert_eq!(declined.operation_id, "second"); + assert_eq!(declined.reason, execute_declined::Reason::AtCapacity as i32); + assert_eq!(declined.detail, "1 of 1 in flight"); + Ok(()) +} + +/// A single-use worker is spent by the action it admits, not by one it +/// turns away: declined for load, it is still free to take the next +/// dispatch. The reservation is read from the property the scheduler named +/// in the connection result. +#[nativelink_test] +async fn a_single_use_worker_is_not_spent_by_a_decline() -> Result<(), Error> { + // A refusal for load needs a free-memory reading, which only Linux has. + if free_memory_kb().is_none() { + return Ok(()); + } + let mut test_context = setup_local_worker_with_config(LocalWorkerConfig { + single_use: true, + worker_api_endpoint: EndpointConfig { + timeout: Some(10000.), + ..Default::default() + }, + ..Default::default() + }) + .await; + let streaming_response = test_context.maybe_streaming_response.take().unwrap(); + test_context + .client + .expect_connect_worker(Ok(streaming_response)) + .await; + + let worker_id = "foobar".to_string(); + let tx_stream = test_context.maybe_tx_stream.take().unwrap(); + tx_stream + .send(Frame::data( + encode_stream_proto(&UpdateForWorker { + update: Some(Update::ConnectionResult(ConnectionResult { + worker_id: worker_id.clone(), + dispatch_ack: true, + memory_property: "memory_kb".to_string(), + })), + }) + .unwrap(), + )) + .await + .map_err(|e| make_input_err!("Could not send : {:?}", e))?; + + let action_info = ActionInfo { + command_digest: DigestInfo::new([1u8; 32], 10), + input_root_digest: DigestInfo::new([2u8; 32], 10), + timeout: Duration::from_secs(1), + platform_properties: HashMap::new(), + priority: 0, + load_timestamp: SystemTime::UNIX_EPOCH, + insert_timestamp: SystemTime::UNIX_EPOCH, + unique_qualifier: ActionUniqueQualifier::Uncacheable(ActionUniqueKey { + execution_scope: None, + instance_name: INSTANCE_NAME.to_string(), + digest_function: DigestHasherFunc::Blake3, + digest: DigestInfo::new([3u8; 32], 10), + }), + }; + let start = |operation_id: &str, memory_kb: Option| UpdateForWorker { + update: Some(Update::StartAction(StartExecute { + request_metadata: None, + execute_request: Some((&action_info).into()), + operation_id: operation_id.to_string(), + queued_timestamp: None, + platform: Some(Platform { + properties: memory_kb + .map(|kb| Property { + name: "memory_kb".to_string(), + value: kb.to_string(), + }) + .into_iter() + .collect(), + }), + worker_id: worker_id.clone(), + })), + }; + + // More memory than any machine has: declined for load. + tx_stream + .send(Frame::data( + encode_stream_proto(&start("greedy", Some(u64::MAX))).unwrap(), + )) + .await + .map_err(|e| make_input_err!("Could not send : {:?}", e))?; + let declined = test_context.client.expect_execute_declined(Ok(())).await; + assert_eq!(declined.operation_id, "greedy"); + assert_eq!(declined.reason, execute_declined::Reason::Load as i32); + assert_eq!(declined.needed_kb, u64::MAX); + + // The worker was not spent by the decline: the next dispatch is admitted. + tx_stream + .send(Frame::data( + encode_stream_proto(&start("modest", None)).unwrap(), + )) + .await + .map_err(|e| make_input_err!("Could not send : {:?}", e))?; + let accepted = test_context.client.expect_execute_accepted(Ok(())).await; + assert_eq!(accepted.operation_id, "modest"); + let running_action = Arc::new(MockRunningAction::new()); + test_context + .actions_manager + .expect_create_and_add_action(Ok(running_action)) + .await; + Ok(()) +} + +/// A precondition script that hangs is refused as backpressure after the +/// timeout instead of holding the action forever. +#[cfg(target_family = "unix")] +#[nativelink_test] +async fn precondition_script_that_hangs_times_out() -> Result<(), Error> { + let extra_envs: HashMap = HashMap::new(); + // A busy loop rather than `sleep`: the script runs with a cleared + // environment, and in a build sandbox only `/bin/sh` is on any path. + let err = preconditions_met( + Some("/bin/sh -c 'while :; do :; done'".to_string()), + &extra_envs, + Duration::from_millis(200), + ) + .await + .expect_err("a hanging script must not pass"); + assert_eq!(err.code, Code::ResourceExhausted, "{err}"); + assert!(err.to_string().contains("did not finish"), "{err}"); + Ok(()) +} + +/// The readiness flag follows the registration: off until the scheduler's +/// `ConnectionResult`, on after it, off again when the connection is lost. +/// The lost state reads Initializing, not Failed, so only the readiness +/// check drops and a liveness probe on the plain status stays green. +#[nativelink_test] +async fn registration_flag_follows_the_scheduler_connection() -> Result<(), Error> { + let mut test_context = setup_local_worker(HashMap::new()).await; + assert!(!test_context.registration.is_registered()); + let streaming_response = test_context.maybe_streaming_response.take().unwrap(); + test_context + .client + .expect_connect_worker(Ok(streaming_response)) + .await; + assert!(!test_context.registration.is_registered()); + + let tx_stream = test_context.maybe_tx_stream.take().unwrap(); + tx_stream + .send(Frame::data( + encode_stream_proto(&UpdateForWorker { + update: Some(Update::ConnectionResult(ConnectionResult { + worker_id: "foobar".to_string(), + dispatch_ack: false, + memory_property: String::new(), + })), + }) + .unwrap(), + )) + .await + .map_err(|e| make_input_err!("Could not send : {:?}", e))?; + let mut registered = false; + for _ in 0..1_000 { + if test_context.registration.is_registered() { + registered = true; + break; + } + tokio::task::yield_now().await; + } + assert!(registered, "registration flag never turned on"); + + drop(tx_stream); + test_context.actions_manager.expect_kill_all().await; + let mut lost = false; + for _ in 0..1_000 { + if !test_context.registration.is_registered() { + lost = true; + break; + } + tokio::task::yield_now().await; + } + assert!( + lost, + "registration flag never turned off after the disconnect" + ); + match test_context.registration.check_health("".into()).await { + HealthStatus::Initializing { message, .. } => { + assert!(message.contains("reconnecting"), "{message}"); + } + other => panic!("a lost registration must read Initializing, got {other:?}"), + } + Ok(()) +} diff --git a/nativelink-worker/tests/namespace_utils_test.rs b/nativelink-worker/tests/namespace_utils_test.rs index b787f952a9b..5648f7f8460 100644 --- a/nativelink-worker/tests/namespace_utils_test.rs +++ b/nativelink-worker/tests/namespace_utils_test.rs @@ -14,10 +14,12 @@ #![cfg(target_os = "linux")] +use core::sync::atomic::{AtomicU64, Ordering}; use core::time::Duration; use std::ffi::CString; use std::os::unix::process::{CommandExt, ExitStatusExt}; -use std::process::Command; +use std::path::{Path, PathBuf}; +use std::process::{Command, Stdio}; use std::time::{SystemTime, UNIX_EPOCH}; use nativelink_error::{Error, ResultExt}; @@ -25,18 +27,253 @@ use nativelink_macro::nativelink_test; use nativelink_worker::namespace_utils; use pretty_assertions::assert_eq; +/// Returns a path under `dir` that is unique for this test run. +fn unique_path(dir: &Path, prefix: &str) -> PathBuf { + static COUNTER: AtomicU64 = AtomicU64::new(0); + #[allow(clippy::cast_possible_truncation)] + let nanos: u64 = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() as u64; + let count = COUNTER.fetch_add(1, Ordering::Relaxed); + dir.join(format!("{prefix}_{}_{nanos}_{count}", std::process::id())) +} + +/// Builds `sh -c script` running in a fresh mount namespace for +/// `action_path` under `root_path`, with a private `/tmp` when +/// `isolate_tmp` is set. +/// Like the worker: with `isolate_tmp` the action's own `tmp` directory is +/// created first and bound over `/tmp` in the child. +fn namespaced_sh(script: &str, isolate_tmp: bool, root_path: &Path, action_path: &Path) -> Command { + let root_dir_c = CString::new(root_path.to_str().unwrap()).unwrap(); + let action_dir_c = CString::new(action_path.to_str().unwrap()).unwrap(); + let tmp_dir_c = isolate_tmp.then(|| { + let tmp_path = action_path.join("tmp"); + std::fs::create_dir_all(&tmp_path).unwrap(); + CString::new(tmp_path.to_str().unwrap()).unwrap() + }); + let mut command = Command::new("sh"); + command.args(["-c", script]); + // SAFETY: configure_namespace is async-signal-safe and intended for pre_exec. + unsafe { + command.pre_exec(move || { + namespace_utils::configure_namespace( + true, + tmp_dir_c.as_deref(), + &root_dir_c, + &action_dir_c, + ) + }); + } + command +} + #[nativelink_test] async fn test_namespaces_supported() -> Result<(), Error> { // This test is a smoke test to ensure that the namespace detection logic // runs without crashing. The result of this function is dependent on the // environment it is run in, so we don't assert the result. - let _supported = namespace_utils::namespaces_supported(false); + let _supported = namespace_utils::namespaces_supported(false, false); + // Isolating /tmp is only possible inside a mount namespace, regardless of + // what the host supports. + assert!(!namespace_utils::namespaces_supported(false, true)); + Ok(()) +} + +#[nativelink_test] +async fn test_configure_namespace_isolate_tmp_hides_host_tmp() -> Result<(), Error> { + if !namespace_utils::namespaces_supported(true, true) { + return Ok(()); + } + + let root_path = unique_path(&std::env::temp_dir(), "nativelink_test_root"); + let action_path = root_path.join("action"); + std::fs::create_dir_all(&action_path).err_tip(|| "Failed to create action dir")?; + + let host_tmp_file = unique_path(Path::new("/tmp"), "nativelink_test_host"); + std::fs::write(&host_tmp_file, "host").err_tip(|| "Failed to write host /tmp file")?; + let scratch_file = unique_path(Path::new("/tmp"), "nativelink_test_scratch"); + + // The host's /tmp contents must be invisible and the private /tmp must be + // writable. + let output = namespaced_sh( + &format!( + "test ! -e {host} && echo scratch > {scratch} && test -f {scratch}", + host = host_tmp_file.display(), + scratch = scratch_file.display(), + ), + true, + &root_path, + &action_path, + ) + .output()?; + std::fs::remove_file(&host_tmp_file).err_tip(|| "Failed to remove host /tmp file")?; + assert_eq!( + Some(0), + output.status.code(), + "Host /tmp was visible or the private /tmp was not writable: {output:?}", + ); + assert!( + !scratch_file.exists(), + "Write to the private /tmp leaked to the host at {}", + scratch_file.display() + ); + // It landed in the action's own tmp directory, which goes with the action. + let landed = action_path + .join("tmp") + .join(scratch_file.file_name().unwrap()); + assert_eq!( + "scratch\n", + std::fs::read_to_string(&landed).err_tip(|| "Failed to read the private /tmp file")?, + "The private /tmp is not the action's tmp directory" + ); + + Ok(()) +} + +#[nativelink_test] +async fn test_configure_namespace_isolate_tmp_keeps_action_directory_under_tmp() -> Result<(), Error> +{ + if !namespace_utils::namespaces_supported(true, true) { + return Ok(()); + } + + // Mirror a worker whose work_directory lives under /tmp, which the private + // tmpfs would otherwise hide. + let root_path = unique_path(Path::new("/tmp"), "nativelink_test_root"); + let action_path = root_path.join("action"); + let work_path = action_path.join("work"); + let sibling_path = root_path.join("sibling"); + let secret_file = sibling_path.join("secret.txt"); + std::fs::create_dir_all(&work_path).err_tip(|| "Failed to create work dir")?; + std::fs::create_dir_all(&sibling_path).err_tip(|| "Failed to create sibling dir")?; + std::fs::write(work_path.join("input.txt"), "input").err_tip(|| "Failed to write input")?; + std::fs::write(&secret_file, "top secret").err_tip(|| "Failed to write secret file")?; + let output_file = work_path.join("output.txt"); + + // Like the worker, start the command inside the action's work directory. + // Inputs must be reachable both relatively and by absolute path, siblings + // must stay masked, and outputs must land in the real work directory. + let mut command = namespaced_sh( + &format!( + "test \"$(pwd)\" = {work} && test -f input.txt && test -f {work}/input.txt && test ! -e {secret} && echo done > output.txt", + work = work_path.display(), + secret = secret_file.display(), + ), + true, + &root_path, + &action_path, + ); + command.current_dir(&work_path); + let output = command.output()?; + assert_eq!( + Some(0), + output.status.code(), + "Action directory under /tmp was not usable: {output:?}", + ); + assert_eq!( + "done\n", + std::fs::read_to_string(&output_file).err_tip(|| "Failed to read output")?, + "Output written in the work directory did not reach the host" + ); + + Ok(()) +} + +#[nativelink_test] +async fn test_configure_namespace_isolate_tmp_concurrent_actions_do_not_collide() +-> Result<(), Error> { + if !namespace_utils::namespaces_supported(true, true) { + return Ok(()); + } + + let root_path = unique_path(&std::env::temp_dir(), "nativelink_test_root"); + let action_path_a = root_path.join("action_a"); + let action_path_b = root_path.join("action_b"); + std::fs::create_dir_all(&action_path_a).err_tip(|| "Failed to create action_a dir")?; + std::fs::create_dir_all(&action_path_b).err_tip(|| "Failed to create action_b dir")?; + + // Both actions use the same predictable path, like a tool that keys its + // scratch file on a pid and assumes it owns /tmp. + let shared_file = unique_path(Path::new("/tmp"), "nativelink_test_shared"); + let script = |owner: &str, delay: &str| { + format!( + "sleep {delay} && echo {owner} > {shared} && sleep 0.5 && cat {shared}", + shared = shared_file.display(), + ) + }; + let mut command_a = namespaced_sh(&script("A", "0"), true, &root_path, &action_path_a); + command_a.stdout(Stdio::piped()).stderr(Stdio::piped()); + let mut command_b = namespaced_sh(&script("B", "0.2"), true, &root_path, &action_path_b); + command_b.stdout(Stdio::piped()).stderr(Stdio::piped()); + + let child_a = command_a.spawn()?; + let child_b = command_b.spawn()?; + let output_a = child_a.wait_with_output()?; + let output_b = child_b.wait_with_output()?; + assert_eq!( + Some(0), + output_a.status.code(), + "Action A failed: {output_a:?}" + ); + assert_eq!( + Some(0), + output_b.status.code(), + "Action B failed: {output_b:?}" + ); + assert_eq!( + "A", + String::from_utf8_lossy(&output_a.stdout).trim(), + "Action A saw action B's write to the shared /tmp path" + ); + assert_eq!( + "B", + String::from_utf8_lossy(&output_b.stdout).trim(), + "Action B saw action A's write to the shared /tmp path" + ); + assert!( + !shared_file.exists(), + "Shared /tmp path leaked to the host at {}", + shared_file.display() + ); + + Ok(()) +} + +#[nativelink_test] +async fn test_configure_namespace_mount_without_isolate_tmp_keeps_host_tmp() -> Result<(), Error> { + if !namespace_utils::namespaces_supported(true, false) { + return Ok(()); + } + + let root_path = unique_path(&std::env::temp_dir(), "nativelink_test_root"); + let action_path = root_path.join("action"); + std::fs::create_dir_all(&action_path).err_tip(|| "Failed to create action dir")?; + + let host_tmp_file = unique_path(Path::new("/tmp"), "nativelink_test_host"); + std::fs::write(&host_tmp_file, "host").err_tip(|| "Failed to write host /tmp file")?; + + // Opting out must leave the host's /tmp visible. + let output = namespaced_sh( + &format!("test -f {}", host_tmp_file.display()), + false, + &root_path, + &action_path, + ) + .output()?; + std::fs::remove_file(&host_tmp_file).err_tip(|| "Failed to remove host /tmp file")?; + assert_eq!( + Some(0), + output.status.code(), + "Host /tmp should stay visible when isolate_tmp is off: {output:?}", + ); + Ok(()) } #[nativelink_test] async fn test_configure_namespace() -> Result<(), Error> { - if !namespace_utils::namespaces_supported(false) { + if !namespace_utils::namespaces_supported(false, false) { return Ok(()); } @@ -50,8 +287,9 @@ async fn test_configure_namespace() -> Result<(), Error> { // It is async-signal-safe and will fork, configure the namespace in the // child, and the original child process will continue to execute the command. unsafe { - command - .pre_exec(move || namespace_utils::configure_namespace(false, &root_dir, &action_dir)); + command.pre_exec(move || { + namespace_utils::configure_namespace(false, None, &root_dir, &action_dir) + }); } let output = command.output()?; @@ -73,7 +311,7 @@ async fn test_configure_namespace() -> Result<(), Error> { #[nativelink_test] async fn test_configure_namespace_mount_isolation() -> Result<(), Error> { - if !namespace_utils::namespaces_supported(true) { + if !namespace_utils::namespaces_supported(true, false) { return Ok(()); } @@ -109,7 +347,7 @@ async fn test_configure_namespace_mount_isolation() -> Result<(), Error> { unsafe { command.pre_exec(move || { - namespace_utils::configure_namespace(true, &root_dir_c, &action1_dir_c) + namespace_utils::configure_namespace(true, None, &root_dir_c, &action1_dir_c) }); } @@ -127,7 +365,7 @@ async fn test_configure_namespace_mount_isolation() -> Result<(), Error> { let action1_dir_c = CString::new(action1_path.to_str().unwrap()).unwrap(); unsafe { command_access.pre_exec(move || { - namespace_utils::configure_namespace(true, &root_dir_c, &action1_dir_c) + namespace_utils::configure_namespace(true, None, &root_dir_c, &action1_dir_c) }); } let output_access = command_access.output()?; @@ -140,9 +378,68 @@ async fn test_configure_namespace_mount_isolation() -> Result<(), Error> { Ok(()) } +/// A SIGTERM to the action's process group reaches the action inside its +/// namespaces as a SIGTERM, so a trap runs and the action can write its +/// cleanup before the SIGKILL that follows the grace. +#[nativelink_test] +async fn test_namespaced_action_gets_sigterm_not_sigkill() -> Result<(), Error> { + if !namespace_utils::namespaces_supported(false, false) { + return Ok(()); + } + let marker = std::env::temp_dir().join(format!( + "nativelink_test_term.{}.{}", + std::process::id(), + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |d| d.as_nanos()) + )); + let mut command = tokio::process::Command::new("sh"); + command.args([ + "-c", + &format!( + "trap 'echo bye > {marker}; exit 0' TERM; sleep 30 & wait", + marker = marker.display() + ), + ]); + // As the worker does: the action leads its own process group, and the + // grace is a SIGTERM to that group. + command.process_group(0); + let root_dir = CString::new("/tmp").unwrap(); + let action_dir = CString::new("/tmp/action").unwrap(); + // SAFETY: configure_namespace is async-signal-safe and intended for pre_exec. + unsafe { + command.pre_exec(move || { + namespace_utils::configure_namespace(false, None, &root_dir, &action_dir) + }); + } + let child = command.spawn()?; + let pgid = i32::try_from(child.id().unwrap()).unwrap(); + let mut namespaced_child = namespace_utils::MaybeNamespacedChild::new(true, child); + // Let the shell install its trap and fork the sleep. + tokio::time::sleep(Duration::from_millis(300)).await; + + // SAFETY: pgid is the group we just created. + assert_eq!(unsafe { libc::killpg(pgid, libc::SIGTERM) }, 0); + let status = tokio::time::timeout(Duration::from_secs(5), namespaced_child.wait()) + .await + .expect("the action should exit on SIGTERM within the grace")?; + assert_eq!( + "bye\n", + std::fs::read_to_string(&marker).err_tip(|| "The trap did not run")?, + "the action did not get to run its TERM trap" + ); + assert_eq!( + status.code(), + Some(0), + "the trap's exit code should come through: {status:?}" + ); + drop(std::fs::remove_file(&marker)); + Ok(()) +} + #[nativelink_test] async fn test_maybe_namespaced_child_kill_reaps_orphans() -> Result<(), Error> { - if !namespace_utils::namespaces_supported(false) { + if !namespace_utils::namespaces_supported(false, false) { return Ok(()); } @@ -161,8 +458,9 @@ async fn test_maybe_namespaced_child_kill_reaps_orphans() -> Result<(), Error> { // SAFETY: configure_namespace is async-signal-safe and intended for pre_exec. unsafe { - command - .pre_exec(move || namespace_utils::configure_namespace(false, &root_dir, &action_dir)); + command.pre_exec(move || { + namespace_utils::configure_namespace(false, None, &root_dir, &action_dir) + }); } let child = command.spawn()?; @@ -171,7 +469,8 @@ async fn test_maybe_namespaced_child_kill_reaps_orphans() -> Result<(), Error> { // Give the shell time to fork the background sleep processes. tokio::time::sleep(Duration::from_millis(200)).await; - // Kill the stub. This sends SIGTERM to the stub, which sends SIGKILL to 'sh'. + // Kill the stub: SIGKILL to the namespace's init takes 'sh' and its + // children with it. namespaced_child.kill().await?; // Ensure the stub process has exited. @@ -228,7 +527,7 @@ async fn test_maybe_namespaced_child_non_namespaced_kill() -> Result<(), Error> #[nativelink_test] async fn test_maybe_namespaced_child_namespaced_natural_exit() -> Result<(), Error> { - if !namespace_utils::namespaces_supported(false) { + if !namespace_utils::namespaces_supported(false, false) { return Ok(()); } @@ -241,8 +540,9 @@ async fn test_maybe_namespaced_child_namespaced_natural_exit() -> Result<(), Err // SAFETY: configure_namespace is async-signal-safe and intended for pre_exec. unsafe { - command - .pre_exec(move || namespace_utils::configure_namespace(false, &root_dir, &action_dir)); + command.pre_exec(move || { + namespace_utils::configure_namespace(false, None, &root_dir, &action_dir) + }); } let child = command.spawn()?; @@ -259,7 +559,7 @@ async fn test_maybe_namespaced_child_namespaced_natural_exit() -> Result<(), Err #[nativelink_test] async fn test_maybe_namespaced_child_try_wait() -> Result<(), Error> { - if !namespace_utils::namespaces_supported(false) { + if !namespace_utils::namespaces_supported(false, false) { return Ok(()); } @@ -269,8 +569,9 @@ async fn test_maybe_namespaced_child_try_wait() -> Result<(), Error> { let root_dir = CString::new("/tmp").unwrap(); let action_dir = CString::new("/tmp/action").unwrap(); unsafe { - command_running - .pre_exec(move || namespace_utils::configure_namespace(false, &root_dir, &action_dir)); + command_running.pre_exec(move || { + namespace_utils::configure_namespace(false, None, &root_dir, &action_dir) + }); } let child_running = command_running.spawn()?; let mut namespaced_child_running = @@ -293,8 +594,9 @@ async fn test_maybe_namespaced_child_try_wait() -> Result<(), Error> { let mut command_exited = tokio::process::Command::new("sh"); command_exited.args(["-c", &format!("exit {expected_exit_code}")]); unsafe { - command_exited - .pre_exec(move || namespace_utils::configure_namespace(false, &root_dir, &action_dir)); + command_exited.pre_exec(move || { + namespace_utils::configure_namespace(false, None, &root_dir, &action_dir) + }); } let child_exited = command_exited.spawn()?; let mut namespaced_child_exited = diff --git a/nativelink-worker/tests/reaper_test.rs b/nativelink-worker/tests/reaper_test.rs new file mode 100644 index 00000000000..9b8459c8247 --- /dev/null +++ b/nativelink-worker/tests/reaper_test.rs @@ -0,0 +1,93 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use nativelink_worker::reaper::parse_zombie_ppid; + +/// Only a zombie's parent is reported, read after the final `)` so a comm +/// with spaces or parentheses cannot shift the fields; a live process, +/// whatever its parent, is nobody's to reap. +#[test] +fn parse_zombie_ppid_reads_only_zombies() { + assert_eq!( + parse_zombie_ppid("42 (sleep) Z 7 42 42 0 -1 4194560 0\n"), + Some(7) + ); + assert_eq!(parse_zombie_ppid("42 (a b) c) Z 7 42 42 0\n"), Some(7)); + assert_eq!(parse_zombie_ppid("42 (sleep) S 7 42 42 0\n"), None); + assert_eq!(parse_zombie_ppid("42 (sleep) Z\n"), None); + assert_eq!(parse_zombie_ppid(""), None); +} + +/// A process an action forked and did not wait for is reparented to the +/// worker and becomes a zombie when it exits. It survives one sweep, which +/// only notes it, and is reaped by the next; a child this process spawned +/// and still holds is never reaped, however long it sits, since its exit +/// status belongs to the handle. +#[cfg(target_os = "linux")] +#[nativelink_macro::nativelink_test] +async fn an_orphaned_zombie_is_reaped_on_the_second_sweep_and_an_owned_one_never() +-> Result<(), nativelink_error::Error> { + use std::collections::BTreeSet; + + use nativelink_worker::reaper::{ + OwnedChild, become_subreaper, reap_orphaned_zombies, zombie_children, + }; + + become_subreaper()?; + // The shell backgrounds a short sleep, prints its pid and exits, which + // orphans the sleep onto us; the sleep then exits and waits for a reaper. + let output = tokio::process::Command::new("sh") + .args(["-c", "sleep 0.2 & echo $!"]) + .output() + .await?; + let orphan: u32 = String::from_utf8_lossy(&output.stdout) + .trim() + .parse() + .unwrap(); + // A child of our own that has exited and that nobody has waited on yet: + // a persistent worker between requests looks like this. + let mut ours = tokio::process::Command::new("true").spawn()?; + let ours_pid = ours.id().expect("just spawned"); + let owned = OwnedChild::new(Some(ours_pid)); + tokio::time::sleep(core::time::Duration::from_millis(600)).await; + let zombies = zombie_children(); + assert!( + zombies.contains(&orphan), + "the orphan should be our zombie by now" + ); + assert!( + zombies.contains(&ours_pid), + "our own child should be a zombie too" + ); + + let (reaped, seen) = reap_orphaned_zombies(&BTreeSet::new()); + assert_eq!(reaped, 0, "the first sweep only notes them"); + assert!(seen.contains(&orphan) && seen.contains(&ours_pid)); + let (reaped, seen) = reap_orphaned_zombies(&seen); + assert_eq!(reaped, 1, "the second sweep reaps the orphan alone"); + assert!( + !std::path::Path::new(&format!("/proc/{orphan}")).exists(), + "the orphan should be gone" + ); + assert!( + !seen.contains(&orphan), + "a reaped pid is not carried into the next sweep" + ); + assert!(seen.contains(&ours_pid), "ours is still there, still noted"); + // The handle collects its own child, as it always could. + let status = ours.wait().await?; + assert!(status.success()); + drop(owned); + Ok(()) +} diff --git a/nativelink-worker/tests/running_actions_manager_test.rs b/nativelink-worker/tests/running_actions_manager_test.rs index 69d6a535c81..b246b65ee61 100644 --- a/nativelink-worker/tests/running_actions_manager_test.rs +++ b/nativelink-worker/tests/running_actions_manager_test.rs @@ -18,6 +18,8 @@ use serial_test::serial; #[serial] mod tests { + #[cfg(target_family = "unix")] + use core::pin::Pin; use core::str::from_utf8; use core::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering}; #[cfg(target_family = "unix")] @@ -29,20 +31,27 @@ mod tests { use std::io::{Cursor, Write}; #[cfg(target_family = "unix")] use std::os::unix::fs::{MetadataExt, OpenOptionsExt}; - use std::path::PathBuf; + use std::path::{Path, PathBuf}; use std::sync::{Arc, LazyLock, Mutex}; use std::time::{SystemTime, UNIX_EPOCH}; + #[cfg(target_family = "unix")] + use async_trait::async_trait; use bytes::Bytes; use futures::prelude::*; use nativelink_config::cas_server::{ - EnvironmentSource, UploadActionResultConfig, UploadCacheResultsStrategy, + DiskEnforcement, EnvironmentSource, MemoryEnforcement, ResourceEnforcementConfig, + UploadActionResultConfig, UploadCacheResultsStrategy, }; use nativelink_config::stores::{ EvictionPolicy, FastSlowSpec, FilesystemSpec, MemorySpec, StoreDirection, StoreSpec, }; use nativelink_error::{Code, Error, ResultExt, make_input_err}; use nativelink_macro::nativelink_test; + #[cfg(target_family = "unix")] + use nativelink_metric::{ + MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent, + }; use nativelink_proto::build::bazel::remote::execution::v2::command::EnvironmentVariable; #[cfg_attr(target_family = "windows", allow(unused_imports))] use nativelink_proto::build::bazel::remote::execution::v2::{ @@ -51,7 +60,7 @@ mod tests { digest_function::Value as ProtoDigestFunction, platform::Property, }; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{ - HistoricalExecuteResponse, StartExecute, + HistoricalExecuteResponse, ResourceOutcome, StartExecute, }; use nativelink_proto::google::rpc::Status; #[cfg(target_family = "unix")] @@ -67,30 +76,50 @@ mod tests { use nativelink_util::action_messages::{ ActionResult, ExecutionMetadata, FileInfo, NameOrPath, OperationId, }; + #[cfg(target_family = "unix")] + use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; use nativelink_util::common::{DigestInfo, fs, make_temp_path}; use nativelink_util::digest_hasher::{DigestHasher, DigestHasherFunc}; + #[cfg(target_family = "unix")] + use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator}; + #[cfg(target_family = "unix")] + use nativelink_util::store_trait::{RemoveItemCallback, StoreDriver, UploadSizeInfo}; use nativelink_util::store_trait::{Store, StoreKey, StoreLike}; use nativelink_worker::directory_cache::{DirectoryCache, DirectoryCacheConfig}; #[cfg(target_os = "linux")] use nativelink_worker::namespace_utils; + #[cfg(target_os = "linux")] + use nativelink_worker::running_actions_manager::MemoryLimit; + #[cfg(target_family = "unix")] + use nativelink_worker::running_actions_manager::free_disk_kb; use nativelink_worker::running_actions_manager::{ - Callbacks, ExecutionConfiguration, RunningAction, RunningActionImpl, RunningActionsManager, - RunningActionsManagerArgs, RunningActionsManagerImpl, download_to_directory, + Callbacks, ExecutionConfiguration, KillReason, PersistentWorkersSettings, + ResourceEnforcement, RunningAction, RunningActionImpl, RunningActionsManager, + RunningActionsManagerArgs, RunningActionsManagerImpl, classify_outcome, + download_to_directory, log_excerpt, parse_kb_field, }; use pretty_assertions::assert_eq; use prost::Message; use tokio::sync::oneshot; + #[cfg(target_family = "unix")] + use tokio::sync::watch; use tracing::info; const DEFAULT_MAX_UPLOAD_TIMEOUT: u64 = 600; + const DEFAULT_MAX_DOWNLOAD_TIMEOUT: u64 = 600; const DEFAULT_MAX_CLEANUP_WAIT: u64 = 30; const DEFAULT_MAX_CLEANUP_BACKOFF: u64 = 500; #[cfg(target_os = "linux")] fn use_namespaces() -> nativelink_worker::running_actions_manager::UseNamespaces { - if namespace_utils::namespaces_supported(true) { - nativelink_worker::running_actions_manager::UseNamespaces::YesAndMount - } else if namespace_utils::namespaces_supported(false) { + // Keep the host's /tmp visible here: several tests stage wrappers and + // payloads under the temp dir, which is /tmp outside of Bazel. The + // private /tmp has its own dedicated test below. + if namespace_utils::namespaces_supported(true, false) { + nativelink_worker::running_actions_manager::UseNamespaces::YesAndMount { + isolate_tmp: false, + } + } else if namespace_utils::namespaces_supported(false, false) { nativelink_worker::running_actions_manager::UseNamespaces::Yes } else { nativelink_worker::running_actions_manager::UseNamespaces::No @@ -349,6 +378,7 @@ mod tests { }, max_action_timeout: Duration::from_mins(10), max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -509,6 +539,7 @@ mod tests { }, max_action_timeout: Duration::from_mins(10), max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -1096,6 +1127,7 @@ mod tests { }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -1225,6 +1257,7 @@ mod tests { }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -1356,6 +1389,7 @@ mod tests { }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -1542,6 +1576,7 @@ mod tests { }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -1730,6 +1765,7 @@ mod tests { }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -1987,6 +2023,7 @@ mod tests { }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -2142,6 +2179,7 @@ mod tests { }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -2288,6 +2326,7 @@ mod tests { }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -2429,6 +2468,7 @@ mod tests { }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -2519,7 +2559,18 @@ mod tests { ) .await?; - let run_action_fut = run_action(running_action_impl); + // The inputs first: a kill that lands while they are still fetching + // ends the action there, and this test is about killing the command. + let running_action_impl = running_action_impl.prepare_action().await?; + let run_action_fut = running_action_impl + .clone() + .execute() + .and_then(RunningAction::upload_results) + .and_then(RunningAction::get_finished_result) + .then(|result| async move { + running_action_impl.cleanup().await?; + result + }); tokio::pin!(run_action_fut); #[cfg(target_family = "unix")] @@ -2626,6 +2677,11 @@ exit 0 Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { root_action_directory: root_action_directory.clone(), execution_configuration: ExecutionConfiguration { + persistent_workers: PersistentWorkersSettings::default(), + max_captured_output_bytes: 0, + resource_enforcement: None, + kill_grace: Duration::ZERO, + set_tmpdir: false, buck2_file_capture: None, entrypoint: Some(test_wrapper_script.into_string().unwrap()), additional_environment: None, @@ -2639,6 +2695,7 @@ exit 0 }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -2792,6 +2849,11 @@ exit 0 Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { root_action_directory: root_action_directory.clone(), execution_configuration: ExecutionConfiguration { + persistent_workers: PersistentWorkersSettings::default(), + max_captured_output_bytes: 0, + resource_enforcement: None, + kill_grace: Duration::ZERO, + set_tmpdir: false, buck2_file_capture: None, entrypoint: Some(test_wrapper_script.into_string().unwrap()), additional_environment: Some(HashMap::from([ @@ -2822,6 +2884,7 @@ exit 0 }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -2926,72 +2989,321 @@ exit 0 Ok(()) } - #[cfg_attr(feature = "nix", ignore)] + /// Each axis of `resource_enforcement` is honoured on its own: the disk + /// guard with memory off, memory with the guard off, and nothing when + /// both are off. #[nativelink_test] - async fn entrypoint_sends_timeout_via_side_channel() -> Result<(), Box> - { - #[cfg(target_family = "unix")] - const TEST_WRAPPER_SCRIPT_CONTENT: &str = "\ -#!/bin/bash -echo '{\"failure\":\"timeout\"}' > \"$SIDE_CHANNEL_FILE\" -exit 1 -"; - #[cfg(target_family = "windows")] - const TEST_WRAPPER_SCRIPT_CONTENT: &str = "\ -@echo off -echo | set /p={\"failure\":\"timeout\"} 1>&2 > %SIDE_CHANNEL_FILE% -exit 1 -"; - const WORKER_ID: &str = "foo_worker_id"; + async fn resource_enforcement_config_honours_each_axis_alone() + -> Result<(), Box> { + let config = |memory, disk| ResourceEnforcementConfig { + memory, + disk, + disk_property_name: "disk_kb".to_string(), + memory_property_name: "memory_kb".to_string(), + memory_headroom_percent: 20, + disk_headroom_percent: 20, + }; + assert!( + ResourceEnforcement::from_config(&config( + MemoryEnforcement::None, + DiskEnforcement::None + )) + .is_none() + ); + let guard_only = ResourceEnforcement::from_config(&config( + MemoryEnforcement::None, + DiskEnforcement::Guard, + )) + .expect("the guard alone is enforcement"); + assert!(guard_only.memory.is_none()); + assert_eq!(guard_only.disk_property_name.as_deref(), Some("disk_kb")); + let memory_only = ResourceEnforcement::from_config(&config( + MemoryEnforcement::Soft, + DiskEnforcement::None, + )) + .expect("memory alone is enforcement"); + let memory = memory_only.memory.expect("the memory axis is on"); + assert_eq!(memory.property_name, "memory_kb"); + assert_eq!(memory.headroom_percent, 20); + assert!(memory_only.disk_property_name.is_none()); + Ok(()) + } - let (_, _, cas_store, ac_store) = setup_stores().await?; + /// The sampler reads `Pss` out of `smaps_rollup`, not a field that + /// happens to share the prefix, and `VmRSS` out of `status`. + #[nativelink_test] + async fn parse_kb_field_reads_the_named_field() -> Result<(), Box> { + let rollup = "00400000-7fff Rollup\nRss: 1234 kB\nPss_Anon: 100 kB\nPss: 567 kB\n"; + assert_eq!(parse_kb_field(rollup, "Pss:"), Some(567)); + assert_eq!(parse_kb_field(rollup, "Rss:"), Some(1234)); + assert_eq!(parse_kb_field("VmRSS:\t 42 kB\n", "VmRSS:"), Some(42)); + assert_eq!(parse_kb_field("VmRSS:\t 42 kB\n", "Pss:"), None); + Ok(()) + } + + /// A manager with only the enforcement given, for the guard tests. Used by the unix-only tests below. + #[cfg(target_family = "unix")] + async fn enforcing_manager( + root_action_directory: &str, + cas_store: &Arc, + execution_configuration: ExecutionConfiguration, + ) -> Result, Error> { + Ok(Arc::new(RunningActionsManagerImpl::new( + RunningActionsManagerArgs { + root_action_directory: root_action_directory.to_string(), + execution_configuration, + cas_store: cas_store.clone(), + ac_store: None, + historical_store: Store::new(cas_store.clone()), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::MAX, + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), + max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + active_input_leases: false, + directory_cache: None, + #[cfg(target_os = "linux")] + use_namespaces: use_namespaces(), + }, + )?)) + } + + /// Uploads a command with an empty input root and returns the action digest. Used by the unix-only tests below. + #[cfg(target_family = "unix")] + async fn upload_action( + cas_store: &Arc, + command: &Command, + ) -> Result { + let command_digest = serialize_and_upload_message( + command, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let input_root_digest = serialize_and_upload_message( + &Directory::default(), + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + serialize_and_upload_message( + &Action { + command_digest: Some(command_digest.into()), + input_root_digest: Some(input_root_digest.into()), + ..Default::default() + }, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await + } + + /// A dispatch of the action carrying the scheduler's properties. Used by the unix-only tests below. + #[cfg(target_family = "unix")] + fn dispatch_with_properties( + action_digest: DigestInfo, + properties: &[(&str, String)], + ) -> StartExecute { + StartExecute { + request_metadata: None, + execute_request: Some(ExecuteRequest { + action_digest: Some(action_digest.into()), + digest_function: ProtoDigestFunction::Sha256.into(), + ..Default::default() + }), + operation_id: OperationId::default().to_string(), + queued_timestamp: None, + platform: Some(Platform { + properties: properties + .iter() + .map(|(name, value)| Property { + name: (*name).to_string(), + value: value.clone(), + }) + .collect(), + }), + worker_id: "foo_worker_id".to_string(), + } + } + + /// The guard counts what the admitted actions reserved: with almost all + /// the free space reserved by one action still running, a second is + /// refused although the disk itself would admit it, and admitted again + /// once the first has cleaned up. + #[cfg(target_family = "unix")] + #[nativelink_test] + async fn disk_guard_counts_what_admitted_actions_reserved() + -> Result<(), Box> { + const WORKER_ID: &str = "foo_worker_id"; + let (_, _, cas_store, _ac_store) = setup_stores().await?; let root_action_directory = make_temp_path("root_action_directory"); fs::create_dir_all(&root_action_directory).await?; - - let test_wrapper_script = { - let test_wrapper_dir = make_temp_path("wrapper_dir"); - fs::create_dir_all(&test_wrapper_dir).await?; - #[cfg(target_family = "unix")] - let test_wrapper_script = OsString::from(test_wrapper_dir + "/test_wrapper_script.sh"); - #[cfg(target_family = "windows")] - let test_wrapper_script = - OsString::from(test_wrapper_dir + "\\test_wrapper_script.bat"); - { - let mut file_options = std::fs::OpenOptions::new(); - file_options.create(true); - file_options.truncate(true); - file_options.write(true); - #[cfg(target_family = "unix")] - file_options.mode(0o777); - let mut test_wrapper_script_handle = file_options - .open(OsString::from(&test_wrapper_script)) - .unwrap(); - test_wrapper_script_handle - .write_all(TEST_WRAPPER_SCRIPT_CONTENT.as_bytes()) - .unwrap(); - test_wrapper_script_handle.sync_all().unwrap(); - // Note: Github runners appear to use some kind of filesystem driver - // that does not sync data as expected. This is the easiest solution. - // See: https://github.com/pantsbuild/pants/issues/10507 - // See: https://github.com/moby/moby/issues/9547 - std::process::Command::new("sync").output().unwrap(); - } - test_wrapper_script + let running_actions_manager = enforcing_manager( + &root_action_directory, + &cas_store, + ExecutionConfiguration { + resource_enforcement: Some(ResourceEnforcement { + memory: None, + disk_property_name: Some("disk_kb".to_string()), + disk_soft: false, + disk_headroom_percent: 20, + }), + ..Default::default() + }, + ) + .await?; + let action_digest = upload_action( + &cas_store, + &Command { + arguments: vec!["true".to_string()], + working_directory: ".".to_string(), + ..Default::default() + }, + ) + .await?; + let free_kb = free_disk_kb(&root_action_directory).expect("statvfs on unix"); + assert!(free_kb > 4096, "the test needs a few MiB free"); + let dispatch = |reserved_kb: u64| { + dispatch_with_properties(action_digest, &[("disk_kb", reserved_kb.to_string())]) }; + // Most of it, admitted and held; the rest is slack for whatever else + // the machine writes while the test runs. + let first = running_actions_manager + .create_and_add_action(WORKER_ID.to_string(), dispatch(free_kb / 4 * 3)) + .await? + .prepare_action() + .await?; + + // Half of what the disk shows: refused, because the first holds it. + let second = running_actions_manager + .create_and_add_action(WORKER_ID.to_string(), dispatch(free_kb / 2)) + .await?; + let err = second + .clone() + .prepare_action() + .await + .expect_err("the free space is spoken for"); + assert_eq!(err.code, Code::ResourceExhausted, "{err}"); + assert!( + err.to_string() + .contains("reserved by the actions already admitted"), + "{err}" + ); + second.cleanup().await?; + + // The first is done; the same reservation now fits. + first.cleanup().await?; + let third = running_actions_manager + .create_and_add_action(WORKER_ID.to_string(), dispatch(free_kb / 2)) + .await? + .prepare_action() + .await?; + third.cleanup().await?; + Ok(()) + } + + /// The scheduler's properties replace the client's only with enforcement + /// on. An environment variable read from a property the client never + /// sent sees the scheduler's value then, and nothing otherwise, so an + /// upgrade alone changes no action's environment. + #[cfg(target_family = "unix")] + #[nativelink_test] + async fn scheduler_properties_reach_the_action_only_with_enforcement_on() + -> Result<(), Box> { + const WORKER_ID: &str = "foo_worker_id"; + let (_, slow_store, cas_store, _ac_store) = setup_stores().await?; + let action_digest = upload_action( + &cas_store, + &Command { + arguments: vec![ + "sh".to_string(), + "-c".to_string(), + "printf '%s' \"$MEM\"".to_string(), + ], + working_directory: ".".to_string(), + environment_variables: vec![EnvironmentVariable { + name: "PATH".to_string(), + value: env::var("PATH").unwrap(), + }], + ..Default::default() + }, + ) + .await?; + let environment = Some(HashMap::from([( + "MEM".to_string(), + EnvironmentSource::Property("memory_kb".to_string()), + )])); + let mut seen = Vec::new(); + for enforcement in [ + Some(ResourceEnforcement { + memory: None, + disk_property_name: Some("disk_kb".to_string()), + disk_soft: false, + disk_headroom_percent: 20, + }), + None, + ] { + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + let running_actions_manager = enforcing_manager( + &root_action_directory, + &cas_store, + ExecutionConfiguration { + resource_enforcement: enforcement, + additional_environment: environment.clone(), + ..Default::default() + }, + ) + .await?; + let running_action = running_actions_manager + .create_and_add_action( + WORKER_ID.to_string(), + dispatch_with_properties(action_digest, &[("memory_kb", "4096".to_string())]), + ) + .await?; + let result = run_action(running_action).await?; + assert_eq!(result.exit_code, 0); + let stdout = slow_store + .as_ref() + .get_part_unchunked(result.stdout_digest, 0, None) + .await?; + seen.push(String::from_utf8(stdout.to_vec())?); + } + assert_eq!(seen, ["4096", ""]); + Ok(()) + } + + #[cfg(target_family = "unix")] + #[nativelink_test] + async fn action_reserving_more_disk_than_is_free_is_refused_as_backpressure() + -> Result<(), Box> { + const WORKER_ID: &str = "foo_worker_id"; + let (_, _, cas_store, _ac_store) = setup_stores().await?; + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + let running_actions_manager = Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { root_action_directory: root_action_directory.clone(), execution_configuration: ExecutionConfiguration { - buck2_file_capture: None, - entrypoint: Some(test_wrapper_script.into_string().unwrap()), - additional_environment: Some(HashMap::from([( - "SIDE_CHANNEL_FILE".to_string(), - EnvironmentSource::SideChannelFile, - )])), + persistent_workers: PersistentWorkersSettings::default(), + // The guard alone: the memory axis is off. + resource_enforcement: Some(ResourceEnforcement { + memory: None, + disk_property_name: Some("disk_kb".to_string()), + disk_soft: false, + disk_headroom_percent: 20, + }), + ..Default::default() }, cas_store: cas_store.clone(), - ac_store: Some(Store::new(ac_store.clone())), + ac_store: None, historical_store: Store::new(cas_store.clone()), upload_action_result_config: &UploadActionResultConfig { upload_ac_results_strategy: UploadCacheResultsStrategy::Never, @@ -2999,6 +3311,7 @@ exit 1 }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -3007,14 +3320,9 @@ exit 1 #[cfg(target_os = "linux")] use_namespaces: use_namespaces(), })?); - let arguments = vec!["true".to_string()]; let command = Command { - arguments, + arguments: vec!["true".to_string()], working_directory: ".".to_string(), - environment_variables: vec![EnvironmentVariable { - name: "PATH".to_string(), - value: env::var("PATH").unwrap(), - }], ..Default::default() }; let command_digest = serialize_and_upload_message( @@ -3041,43 +3349,2152 @@ exit 1 ) .await?; - let execute_request = ExecuteRequest { - action_digest: Some(action_digest.into()), - digest_function: ProtoDigestFunction::Sha256.into(), - ..Default::default() - }; - let operation_id = OperationId::default().to_string(); - let running_action_impl = running_actions_manager - .clone() .create_and_add_action( WORKER_ID.to_string(), StartExecute { request_metadata: None, - execute_request: Some(execute_request), - operation_id, - queued_timestamp: Some(make_system_time(1000).into()), - platform: action.platform.clone(), + execute_request: Some(ExecuteRequest { + action_digest: Some(action_digest.into()), + digest_function: ProtoDigestFunction::Sha256.into(), + ..Default::default() + }), + operation_id: OperationId::default().to_string(), + queued_timestamp: None, + // More than any test machine has free. + platform: Some(Platform { + properties: vec![Property { + name: "disk_kb".into(), + value: (u64::MAX / 2).to_string(), + }], + }), worker_id: WORKER_ID.to_string(), }, ) .await?; - let result = run_action(running_action_impl).await?; - assert_eq!(result.exit_code, 1, "Exit code should be 1"); - assert_eq!( - result.error.err_tip(|| "Error should exist")?.code, - Code::DeadlineExceeded + let err = run_action(running_action_impl) + .await + .expect_err("an action reserving more disk than is free must be refused"); + assert_eq!(err.code, Code::ResourceExhausted, "{err}"); + assert!( + err.to_string().contains("Not enough free disk"), + "the refusal should name the disk: {err}" ); - assert!(logs_contain( - "Command returned non-zero exit code exit_code=1 stdout=\"\" stderr=\"\" command=[\"true\"]" - )); Ok(()) } + #[cfg(target_os = "linux")] #[nativelink_test] - async fn caches_results_in_action_cache_store() -> Result<(), Box> { - let (_, _, cas_store, ac_store) = setup_stores().await?; + async fn timed_out_action_takes_its_children_with_it() -> Result<(), Box> + { + const WORKER_ID: &str = "foo_worker_id"; + let (_, _, cas_store, _ac_store) = setup_stores().await?; + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + // Under the (existing) action root, so its parent directory exists. + let pid_file = format!("{root_action_directory}/children.pids"); + + let running_actions_manager = + Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { + root_action_directory: root_action_directory.clone(), + execution_configuration: ExecutionConfiguration::default(), + cas_store: cas_store.clone(), + ac_store: None, + historical_store: Store::new(cas_store.clone()), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::MAX, + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), + max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + active_input_leases: false, + directory_cache: None, + // Host pids and a host-visible pid file: the point is what the kill + // reaches, not the namespaces. + use_namespaces: nativelink_worker::running_actions_manager::UseNamespaces::No, + })?); + // Two grandchildren that would outlive a kill aimed at the shell + // alone; their pids land in a file outside the action directory. + let command = Command { + arguments: vec![ + "sh".to_string(), + "-c".to_string(), + format!("sleep 60 & echo $! > {pid_file}; sleep 60 & echo $! >> {pid_file}; wait"), + ], + working_directory: ".".to_string(), + environment_variables: vec![EnvironmentVariable { + name: "PATH".to_string(), + value: env::var("PATH").unwrap(), + }], + ..Default::default() + }; + let command_digest = serialize_and_upload_message( + &command, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let input_root_digest = serialize_and_upload_message( + &Directory::default(), + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let action = Action { + command_digest: Some(command_digest.into()), + input_root_digest: Some(input_root_digest.into()), + timeout: Some(prost_types::Duration { + seconds: 1, + nanos: 0, + }), + ..Default::default() + }; + let action_digest = serialize_and_upload_message( + &action, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + + let running_action_impl = running_actions_manager + .create_and_add_action( + WORKER_ID.to_string(), + StartExecute { + request_metadata: None, + execute_request: Some(ExecuteRequest { + action_digest: Some(action_digest.into()), + digest_function: ProtoDigestFunction::Sha256.into(), + ..Default::default() + }), + operation_id: OperationId::default().to_string(), + queued_timestamp: None, + platform: None, + worker_id: WORKER_ID.to_string(), + }, + ) + .await?; + + let result = run_action(running_action_impl).await?; + let err = result + .error + .expect("a timed-out action must carry an error"); + assert_eq!(err.code, Code::DeadlineExceeded, "{err}"); + + let pids: Vec = std::fs::read_to_string(&pid_file)? + .lines() + .filter_map(|line| line.trim().parse().ok()) + .collect(); + assert_eq!(pids.len(), 2, "both children should have written their pid"); + // A moment for the kernel to reap; then neither pid may exist. + tokio::time::sleep(Duration::from_millis(200)).await; + for pid in pids { + // SAFETY: signal 0 checks for existence and sends nothing. + let alive = unsafe { libc::kill(pid, 0) } == 0; + assert!(!alive, "child {pid} outlived the timed-out action"); + } + Ok(()) + } + + /// A store whose uploads wait at a gate, counting how many are waiting + /// at once. Every upload of an output file arrives here with that file + /// open, so the peak is the number of files the worker had open together. + /// Unix only, with the one test that drives it. + #[cfg(target_family = "unix")] + #[derive(Debug)] + struct GatedStore { + inner: Store, + gate: watch::Sender, + in_flight: AtomicU64, + max_in_flight: AtomicU64, + } + + #[cfg(target_family = "unix")] + impl GatedStore { + fn new(inner: Arc) -> Arc { + Arc::new(Self { + inner: Store::new(inner), + gate: watch::Sender::new(true), + in_flight: AtomicU64::new(0), + max_in_flight: AtomicU64::new(0), + }) + } + + fn set_open(&self, open: bool) { + self.gate.send_replace(open); + } + + fn in_flight(&self) -> u64 { + self.in_flight.load(Ordering::SeqCst) + } + + fn max_in_flight(&self) -> u64 { + self.max_in_flight.load(Ordering::SeqCst) + } + } + + #[cfg(target_family = "unix")] + impl MetricsComponent for GatedStore { + fn publish( + &self, + _kind: MetricKind, + _field_metadata: MetricFieldData, + ) -> Result { + Ok(MetricPublishKnownKindData::Component) + } + } + + #[cfg(target_family = "unix")] + #[async_trait] + impl StoreDriver for GatedStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + + async fn has_with_results( + self: Pin<&Self>, + keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + self.inner + .as_store_driver_pin() + .has_with_results(keys, results) + .await + } + + async fn update( + self: Pin<&Self>, + key: StoreKey<'_>, + reader: DropCloserReadHalf, + size_info: UploadSizeInfo, + ) -> Result { + let now_in_flight = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; + self.max_in_flight + .fetch_max(now_in_flight, Ordering::SeqCst); + let mut gate = self.gate.subscribe(); + gate.wait_for(|open| *open) + .await + .map_err(|_| make_input_err!("gate dropped"))?; + let result = self + .inner + .as_store_driver_pin() + .update(key, reader, size_info) + .await; + self.in_flight.fetch_sub(1, Ordering::SeqCst); + result + } + + async fn get_part( + self: Pin<&Self>, + key: StoreKey<'_>, + writer: &mut DropCloserWriteHalf, + offset: u64, + length: Option, + ) -> Result<(), Error> { + self.inner + .as_store_driver_pin() + .get_part(key, writer, offset, length) + .await + } + + fn inner_store(&self, _key: Option) -> &dyn StoreDriver { + self + } + + fn as_any<'a>(&'a self) -> &'a (dyn core::any::Any + Sync + Send + 'static) { + self + } + + fn as_any_arc(self: Arc) -> Arc { + self + } + + fn register_remove_callback( + self: Arc, + _callback: Arc, + ) -> Result<(), Error> { + Ok(()) + } + } + + #[cfg(target_family = "unix")] + default_health_status_indicator!(GatedStore); + + /// A store that answers `has` and uploads normally but can be told to + /// never finish a read of one digest, or any write: what a wedged tier + /// looks like to the worker. + #[cfg(target_family = "unix")] + #[derive(Debug)] + struct HangingStore { + inner: Store, + hang_on: DigestInfo, + hang_reads: AtomicBool, + hang_updates: AtomicBool, + } + + #[cfg(target_family = "unix")] + impl HangingStore { + fn new(inner: Arc, hang_on: DigestInfo) -> Arc { + Arc::new(Self { + inner: Store::new(inner), + hang_on, + hang_reads: AtomicBool::new(true), + hang_updates: AtomicBool::new(false), + }) + } + } + + #[cfg(target_family = "unix")] + impl MetricsComponent for HangingStore { + fn publish( + &self, + _kind: MetricKind, + _field_metadata: MetricFieldData, + ) -> Result { + Ok(MetricPublishKnownKindData::Component) + } + } + + #[cfg(target_family = "unix")] + #[async_trait] + impl StoreDriver for HangingStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + async fn has_with_results( + self: Pin<&Self>, + keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + self.inner + .as_store_driver_pin() + .has_with_results(keys, results) + .await + } + async fn update( + self: Pin<&Self>, + key: StoreKey<'_>, + reader: DropCloserReadHalf, + size_info: UploadSizeInfo, + ) -> Result { + if self.hang_updates.load(Ordering::Acquire) { + future::pending::<()>().await; + } + self.inner + .as_store_driver_pin() + .update(key, reader, size_info) + .await + } + async fn get_part( + self: Pin<&Self>, + key: StoreKey<'_>, + writer: &mut DropCloserWriteHalf, + offset: u64, + length: Option, + ) -> Result<(), Error> { + if self.hang_reads.load(Ordering::Acquire) && key == StoreKey::Digest(self.hang_on) { + future::pending::<()>().await; + } + self.inner + .as_store_driver_pin() + .get_part(key, writer, offset, length) + .await + } + fn inner_store(&self, _key: Option) -> &dyn StoreDriver { + self + } + fn as_any<'a>(&'a self) -> &'a (dyn core::any::Any + Sync + Send + 'static) { + self + } + fn as_any_arc(self: Arc) -> Arc { + self + } + fn register_remove_callback( + self: Arc, + _callback: Arc, + ) -> Result<(), Error> { + Ok(()) + } + } + + #[cfg(target_family = "unix")] + default_health_status_indicator!(HangingStore); + + /// A manager whose CAS never finishes reading one input of the action it + /// returns the digest of; the store comes back too, so a test can lift + /// that or make writes hang instead. + #[cfg(target_family = "unix")] + async fn hung_fetch_setup( + max_download_timeout: Duration, + ) -> Result< + ( + Arc, + DigestInfo, + Arc, + ), + Box, + > { + const FILE_CONTENT: &[u8] = b"never arrives"; + let inner = MemoryStore::new(&MemorySpec::default()); + let file_digest = DigestInfo::new([9u8; 32], FILE_CONTENT.len() as u64); + inner + .as_ref() + .update_oneshot(file_digest, FILE_CONTENT.into()) + .await?; + let slow_store = HangingStore::new(inner.clone(), file_digest); + let fast_config = FilesystemSpec { + content_path: make_temp_path("content_path"), + temp_path: make_temp_path("temp_path"), + eviction_policy: None, + ..Default::default() + }; + let fast_store: Arc = FilesystemStore::new(&fast_config).await?; + let cas_store = FastSlowStore::new( + &FastSlowSpec { + fast: StoreSpec::Filesystem(fast_config), + slow: StoreSpec::Memory(MemorySpec::default()), + fast_direction: StoreDirection::default(), + slow_direction: StoreDirection::default(), + bypass_dedup_threshold_bytes: 0, + }, + Store::new(fast_store), + Store::new(slow_store.clone()), + ); + let input_root_digest = serialize_and_upload_message( + &Directory { + files: vec![FileNode { + name: "input.txt".to_string(), + digest: Some(file_digest.into()), + ..Default::default() + }], + ..Default::default() + }, + inner.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let command_digest = serialize_and_upload_message( + &Command { + arguments: vec!["sh".to_string(), "-c".to_string(), "echo hi".to_string()], + ..Default::default() + }, + inner.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let action_digest = serialize_and_upload_message( + &Action { + command_digest: Some(command_digest.into()), + input_root_digest: Some(input_root_digest.into()), + ..Default::default() + }, + inner.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + let running_actions_manager = + Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { + root_action_directory, + execution_configuration: ExecutionConfiguration::default(), + cas_store: cas_store.clone(), + ac_store: None, + historical_store: Store::new(cas_store.clone()), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::from_mins(10), + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout, + max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + active_input_leases: false, + directory_cache: None, + #[cfg(target_os = "linux")] + use_namespaces: use_namespaces(), + })?); + Ok((running_actions_manager, action_digest, slow_store)) + } + + #[cfg(target_family = "unix")] + fn start_execute_for(action_digest: DigestInfo, operation_id: &str) -> StartExecute { + StartExecute { + request_metadata: None, + execute_request: Some(ExecuteRequest { + action_digest: Some(action_digest.into()), + digest_function: ProtoDigestFunction::Sha256.into(), + ..Default::default() + }), + operation_id: operation_id.to_string(), + ..Default::default() + } + } + + /// An input whose fetch never completes fails the action at + /// `max_download_timeout` with `DeadlineExceeded` instead of holding its + /// slot for good, and the action then cleans up like any other. + #[cfg(target_family = "unix")] + #[nativelink_test] + async fn prepare_action_times_out_on_a_hung_fetch() -> Result<(), Box> { + let (running_actions_manager, action_digest, _store) = + hung_fetch_setup(Duration::from_secs(1)).await?; + let operation_id = OperationId::default(); + let started = std::time::Instant::now(); + let action = running_actions_manager + .create_and_add_action( + "test-worker".to_string(), + start_execute_for(action_digest, &operation_id.to_string()), + ) + .await?; + let action_directory = action.get_work_directory().clone(); + let result = action.clone().prepare_action().await; + let err = result.expect_err("the fetch never completes, prepare must fail"); + assert_eq!(err.code, Code::DeadlineExceeded, "{err:?}"); + assert!(err.to_string().contains("max_download_timeout"), "{err:?}"); + assert!( + started.elapsed() < Duration::from_secs(30), + "took {:?}", + started.elapsed() + ); + action.cleanup().await?; + assert!( + !Path::new(&action_directory).exists(), + "the work directory is gone after cleanup" + ); + tokio::time::timeout(Duration::from_secs(5), running_actions_manager.kill_all()) + .await + .expect("an entry was left behind"); + Ok(()) + } + + /// A kill that arrives while the inputs are still fetching ends the + /// action at once with `Aborted`; it does not wait for the fetch, a + /// later `execute` fails instead of panicking, and cleanup runs. + #[cfg(target_family = "unix")] + #[nativelink_test] + async fn prepare_action_stops_on_a_kill() -> Result<(), Box> { + let (running_actions_manager, action_digest, _store) = + hung_fetch_setup(Duration::from_mins(1)).await?; + let operation_id = OperationId::default(); + let action = running_actions_manager + .create_and_add_action( + "test-worker".to_string(), + start_execute_for(action_digest, &operation_id.to_string()), + ) + .await?; + let action_directory = action.get_work_directory().clone(); + let started = std::time::Instant::now(); + let killer = { + let manager = running_actions_manager.clone(); + let operation_id = operation_id.clone(); + async move { + tokio::time::sleep(Duration::from_millis(200)).await; + manager.kill_operation(&operation_id).await + } + }; + let (kill_res, prepare_res) = tokio::join!(killer, action.clone().prepare_action()); + kill_res?; + let err = prepare_res.expect_err("a killed fetch must not complete"); + assert_eq!(err.code, Code::Aborted, "{err:?}"); + assert!( + started.elapsed() < Duration::from_secs(30), + "took {:?}", + started.elapsed() + ); + let err = action + .clone() + .execute() + .await + .expect_err("execute after a killed prepare must fail"); + assert_ne!(err.code, Code::Ok, "{err:?}"); + action.cleanup().await?; + assert!( + !Path::new(&action_directory).exists(), + "the work directory is gone after cleanup" + ); + Ok(()) + } + + /// A kill that arrives while the results upload ends it at once with + /// `Aborted`: a kill is acknowledged in every phase, not only while the + /// command runs. + #[cfg(target_family = "unix")] + #[nativelink_test] + async fn upload_results_stops_on_a_kill() -> Result<(), Box> { + let (running_actions_manager, action_digest, store) = + hung_fetch_setup(Duration::from_mins(1)).await?; + // The fetch completes; the write of the results never does. + store.hang_reads.store(false, Ordering::Release); + store.hang_updates.store(true, Ordering::Release); + let operation_id = OperationId::default(); + let action = running_actions_manager + .create_and_add_action( + "test-worker".to_string(), + start_execute_for(action_digest, &operation_id.to_string()), + ) + .await?; + let action_directory = action.get_work_directory().clone(); + let started = std::time::Instant::now(); + let killer = { + let manager = running_actions_manager.clone(); + let operation_id = operation_id.clone(); + async move { + tokio::time::sleep(Duration::from_millis(500)).await; + manager.kill_operation(&operation_id).await + } + }; + let (kill_res, upload_res) = tokio::join!(killer, async { + action + .clone() + .prepare_action() + .await? + .execute() + .await? + .upload_results() + .await + }); + kill_res?; + let err = upload_res.expect_err("a killed upload must not complete"); + assert_eq!(err.code, Code::Aborted, "{err:?}"); + assert!( + started.elapsed() < Duration::from_secs(30), + "took {:?}", + started.elapsed() + ); + action.cleanup().await?; + assert!(!Path::new(&action_directory).exists()); + Ok(()) + } + + /// The bound on open output files holds across the whole tree, not per + /// directory: eight directories of thirty-two files each are uploaded + /// through one gate, and never more than the bound (plus stdout and + /// stderr, which hold no file) are in flight at once. Per-directory + /// bounds would have all 256 open together. + #[cfg(target_family = "unix")] + #[nativelink_test] + async fn output_uploads_share_one_open_file_bound() -> Result<(), Box> { + const WORKER_ID: &str = "foo_worker_id"; + const DIRS: u64 = 8; + const FILES_PER_DIR: u64 = 32; + const BOUND: u64 = 64; + // The gate is the slow store: the manager insists on a filesystem + // fast store, and an upload is not done until both sides have it. + let fast_config = FilesystemSpec { + content_path: make_temp_path("content_path"), + temp_path: make_temp_path("temp_path"), + eviction_policy: None, + ..Default::default() + }; + let fast_store: Arc = FilesystemStore::new(&fast_config).await?; + let gated = GatedStore::new(MemoryStore::new(&MemorySpec::default())); + let cas_store = FastSlowStore::new( + &FastSlowSpec { + fast: StoreSpec::Filesystem(fast_config), + slow: StoreSpec::Memory(MemorySpec::default()), + fast_direction: StoreDirection::default(), + slow_direction: StoreDirection::default(), + bypass_dedup_threshold_bytes: 0, + }, + Store::new(fast_store), + Store::new(gated.clone()), + ); + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + let running_actions_manager = + Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { + root_action_directory: root_action_directory.clone(), + execution_configuration: ExecutionConfiguration::default(), + cas_store: cas_store.clone(), + ac_store: None, + historical_store: Store::new(cas_store.clone()), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::MAX, + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), + max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + active_input_leases: false, + directory_cache: None, + #[cfg(target_os = "linux")] + use_namespaces: use_namespaces(), + })?); + // Every file's content is unique, so none is skipped as a duplicate. + let command = Command { + arguments: vec![ + "sh".to_string(), + "-c".to_string(), + format!( + "for d in $(seq {DIRS}); do mkdir -p out/d$d; for f in $(seq {FILES_PER_DIR}); do echo $d-$f > out/d$d/f$f; done; done" + ), + ], + working_directory: ".".to_string(), + environment_variables: vec![EnvironmentVariable { + name: "PATH".to_string(), + value: env::var("PATH").unwrap(), + }], + output_paths: vec!["out".to_string()], + ..Default::default() + }; + let command_digest = serialize_and_upload_message( + &command, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let input_root_digest = serialize_and_upload_message( + &Directory::default(), + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let action = Action { + command_digest: Some(command_digest.into()), + input_root_digest: Some(input_root_digest.into()), + ..Default::default() + }; + let action_digest = serialize_and_upload_message( + &action, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let running_action_impl = running_actions_manager + .create_and_add_action( + WORKER_ID.to_string(), + StartExecute { + request_metadata: None, + execute_request: Some(ExecuteRequest { + action_digest: Some(action_digest.into()), + digest_function: ProtoDigestFunction::Sha256.into(), + ..Default::default() + }), + operation_id: OperationId::default().to_string(), + queued_timestamp: None, + platform: None, + worker_id: WORKER_ID.to_string(), + }, + ) + .await?; + + // Inputs are fetched and the command runs with the gate open; the + // uploads then pile up against it. + let running_action_impl = running_action_impl + .clone() + .prepare_action() + .and_then(RunningAction::execute) + .await?; + gated.set_open(false); + let upload = tokio::spawn(running_action_impl.upload_results()); + let filled = async { + while gated.in_flight() < BOUND { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }; + tokio::time::timeout(Duration::from_secs(30), filled) + .await + .expect("the uploads should reach the bound"); + // Anything more that could be in flight arrives in this window. + tokio::time::sleep(Duration::from_millis(200)).await; + let peak_at_gate = gated.in_flight(); + gated.set_open(true); + let running_action_impl = upload.await??; + let result = running_action_impl.clone().get_finished_result().await?; + running_action_impl.cleanup().await?; + assert_eq!(result.exit_code, 0); + assert_eq!(result.output_folders.len(), 1); + + // stdout and stderr are the two uploads that hold no file. + assert!( + peak_at_gate <= BOUND + 2 && gated.max_in_flight() <= BOUND + 2, + "{} at the gate, {} at most: more files open than the bound", + peak_at_gate, + gated.max_in_flight() + ); + assert!( + gated.max_in_flight() >= BOUND, + "{} in flight at most: the bound was never reached", + gated.max_in_flight() + ); + Ok(()) + } + #[cfg(target_os = "linux")] + #[nativelink_test] + async fn action_writing_past_its_disk_reservation_is_killed() + -> Result<(), Box> { + const WORKER_ID: &str = "foo_worker_id"; + let (_, _, cas_store, _ac_store) = setup_stores().await?; + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + + let running_actions_manager = + Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { + root_action_directory: root_action_directory.clone(), + execution_configuration: ExecutionConfiguration { + persistent_workers: PersistentWorkersSettings::default(), + resource_enforcement: Some(ResourceEnforcement { + memory: None, + disk_property_name: Some("disk_kb".to_string()), + disk_soft: true, + disk_headroom_percent: 20, + }), + ..Default::default() + }, + cas_store: cas_store.clone(), + ac_store: None, + historical_store: Store::new(cas_store.clone()), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::MAX, + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), + max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + active_input_leases: false, + directory_cache: None, + #[cfg(target_os = "linux")] + use_namespaces: use_namespaces(), + })?); + let command = Command { + arguments: vec![ + "sh".to_string(), + "-c".to_string(), + // Thirty MiB of its own under the action directory, then a + // long sleep: a disk sample sees it past the reservation. + "dd if=/dev/zero of=fill.bin bs=1M count=30 status=none; sleep 60".to_string(), + ], + working_directory: ".".to_string(), + environment_variables: vec![EnvironmentVariable { + name: "PATH".to_string(), + value: env::var("PATH").unwrap(), + }], + ..Default::default() + }; + let command_digest = serialize_and_upload_message( + &command, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let input_root_digest = serialize_and_upload_message( + &Directory::default(), + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + // The client tagged nothing; the reservation arrives the way a hint or + // cold start does, on the scheduler's StartExecute, not on the Action. + let action = Action { + command_digest: Some(command_digest.into()), + input_root_digest: Some(input_root_digest.into()), + ..Default::default() + }; + let action_digest = serialize_and_upload_message( + &action, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + + let running_action_impl = running_actions_manager + .create_and_add_action( + WORKER_ID.to_string(), + StartExecute { + request_metadata: None, + execute_request: Some(ExecuteRequest { + action_digest: Some(action_digest.into()), + digest_function: ProtoDigestFunction::Sha256.into(), + ..Default::default() + }), + operation_id: OperationId::default().to_string(), + queued_timestamp: None, + // 1 MiB of disk reserved: thirty MiB written is far past + // it with 20% headroom. + platform: Some(Platform { + properties: vec![Property { + name: "disk_kb".into(), + value: "1024".into(), + }], + }), + worker_id: WORKER_ID.to_string(), + }, + ) + .await?; + + let started = std::time::Instant::now(); + let result = run_action(running_action_impl.clone()).await?; + let err = result + .error + .expect("an action over its reservation must fail with an error"); + assert_eq!(err.code, Code::FailedPrecondition, "{err}"); + assert!( + err.to_string().contains("disk reservation"), + "error should name the reservation: {err}" + ); + assert!( + started.elapsed() < Duration::from_secs(30), + "the worker should have killed the action at a disk sample, long before its sleep ended" + ); + // The measurement says what happened and what it was measured against. + let usage = running_action_impl + .resource_usage() + .expect("usage is reported for every executed action"); + assert_eq!( + usage.outcome, + ResourceOutcome::KilledDisk as i32, + "{usage:?}" + ); + assert!(usage.enforced, "the worker's own limit ended it: {usage:?}"); + // The sample that ended it caught the write in progress, so the + // figure is whatever had landed by then: past the reservation, not + // necessarily the whole thirty MiB. + assert!( + usage.peak_disk_kb > 1024, + "the measurement should carry what the action had written: {usage:?}" + ); + Ok(()) + } + + /// Only what the action wrote counts: files modified at or after the + /// directory's stamp, taken from the filesystem's own clock, which on + /// Linux lags the wall clock by a tick. An input materialized before it, hard-linked or copied and + /// whatever its link count now, keeps its earlier time and is not + /// counted; nor are a symlink or the directories' own blocks. + #[cfg(target_family = "unix")] + #[test] + fn directory_private_kb_counts_only_files_newer_than_the_start() { + let dir = tempfile::tempdir().unwrap(); + let before = SystemTime::now() - Duration::from_secs(90); + // An input with one link, as after the store evicted its blob, and + // one with two: both older than the action, both left out. + let input = dir.path().join("input.bin"); + std::fs::write(&input, vec![0u8; 64 * 1024]).unwrap(); + std::fs::File::open(&input) + .unwrap() + .set_modified(before) + .unwrap(); + let linked = dir.path().join("linked.bin"); + std::fs::write(&linked, vec![0u8; 64 * 1024]).unwrap(); + std::fs::File::open(&linked) + .unwrap() + .set_modified(before) + .unwrap(); + std::fs::hard_link(&linked, dir.path().join("linked-2.bin")).unwrap(); + let since = + nativelink_worker::running_actions_manager::directory_stamp(dir.path()).unwrap(); + let own = dir.path().join("own.bin"); + std::fs::write(&own, vec![0u8; 64 * 1024]).unwrap(); + std::fs::create_dir(dir.path().join("sub")).unwrap(); + std::fs::write( + dir.path().join("sub").join("deep.bin"), + vec![0u8; 32 * 1024], + ) + .unwrap(); + std::os::unix::fs::symlink(&own, dir.path().join("own-symlink")).unwrap(); + let kb = + nativelink_worker::running_actions_manager::directory_private_kb(dir.path(), since); + assert!( + (90..=100).contains(&kb), + "own 64 KiB plus deep 32 KiB, not the inputs: {kb} KiB" + ); + } + + /// A timed-out action gets SIGTERM first and keeps what it printed; + /// SIGKILL follows only when it ignores the grace period. Linux only: + /// the worker signals the process group, which only Linux sets up. + #[cfg(target_os = "linux")] + #[nativelink_test] + async fn timed_out_action_gets_a_grace_period_and_keeps_its_output() + -> Result<(), Box> { + const WORKER_ID: &str = "foo_worker_id"; + let (_, _, cas_store, _ac_store) = setup_stores().await?; + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + let running_actions_manager = + Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { + root_action_directory: root_action_directory.clone(), + execution_configuration: ExecutionConfiguration { + persistent_workers: PersistentWorkersSettings::default(), + kill_grace: Duration::from_secs(5), + ..Default::default() + }, + cas_store: cas_store.clone(), + ac_store: None, + historical_store: Store::new(cas_store.clone()), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::MAX, + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), + max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + active_input_leases: false, + directory_cache: None, + // The namespaced stub turns SIGTERM into SIGKILL for its child, + // so the grace period is only observable without it. + #[cfg(target_os = "linux")] + use_namespaces: nativelink_worker::running_actions_manager::UseNamespaces::No, + })?); + let command = Command { + arguments: vec![ + "sh".to_string(), + "-c".to_string(), + "echo partial; trap 'echo cleaned > cleanup.marker; exit 3' TERM; while :; do sleep 0.1; done" + .to_string(), + ], + working_directory: ".".to_string(), + environment_variables: vec![EnvironmentVariable { + name: "PATH".to_string(), + value: env::var("PATH").unwrap(), + }], + ..Default::default() + }; + let command_digest = serialize_and_upload_message( + &command, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let input_root_digest = serialize_and_upload_message( + &Directory::default(), + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let action = Action { + command_digest: Some(command_digest.into()), + input_root_digest: Some(input_root_digest.into()), + timeout: Some(prost_types::Duration { + seconds: 1, + nanos: 0, + }), + ..Default::default() + }; + let action_digest = serialize_and_upload_message( + &action, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let running_action = running_actions_manager + .create_and_add_action( + WORKER_ID.to_string(), + StartExecute { + request_metadata: None, + execute_request: Some(ExecuteRequest { + action_digest: Some(action_digest.into()), + digest_function: ProtoDigestFunction::Sha256.into(), + ..Default::default() + }), + operation_id: OperationId::default().to_string(), + queued_timestamp: None, + platform: None, + worker_id: WORKER_ID.to_string(), + }, + ) + .await?; + let started = std::time::Instant::now(); + let running_action = running_action.prepare_action().await?.execute().await?; + assert!( + started.elapsed() < Duration::from_secs(4), + "the shell exits on SIGTERM, so the grace period should not run out" + ); + // The trap ran, so SIGTERM arrived before any SIGKILL. + let action_directory = std::fs::read_dir(&root_action_directory)? + .filter_map(Result::ok) + .find(|entry| entry.path().is_dir()) + .expect("the action directory exists until cleanup") + .path(); + assert!( + action_directory.join("work/cleanup.marker").exists(), + "the shell's TERM trap should have written its marker" + ); + let running_action = running_action.upload_results().await?; + let usage = running_action + .resource_usage() + .expect("usage is reported for every executed action"); + assert_eq!( + usage.outcome, + ResourceOutcome::KilledTimeout as i32, + "{usage:?}" + ); + assert!(!usage.enforced, "{usage:?}"); + assert!(usage.wall_time_ms >= 900, "{usage:?}"); + let result = running_action.clone().get_finished_result().await?; + running_action.cleanup().await?; + let err = result + .error + .as_ref() + .expect("a timed-out action carries the error"); + assert_eq!(err.code, Code::DeadlineExceeded, "{err}"); + // The shell left through its trap, not through SIGKILL. + assert_eq!(result.exit_code, 3, "{result:?}"); + // What it printed before the kill was uploaded, not dropped. + let partial = compute_buf_digest(b"partial\n", &mut DigestHasherFunc::Sha256.hasher()); + assert_eq!(result.stdout_digest, partial, "{result:?}"); + assert!( + cas_store + .as_pin() + .has(StoreKey::Digest(partial)) + .await? + .is_some(), + "the killed action's stdout should be in CAS" + ); + Ok(()) + } + + /// An action that exits leaving a process behind holding its stdout + /// finishes when it exits, with what it printed, instead of sitting + /// there until its timeout. On Linux the straggler is killed with the + /// group; elsewhere the drain deadline ends the wait. + #[cfg(target_family = "unix")] + #[nativelink_test] + async fn an_action_that_leaves_a_process_on_its_pipes_still_finishes() + -> Result<(), Box> { + const WORKER_ID: &str = "foo_worker_id"; + let (_, slow_store, cas_store, _ac_store) = setup_stores().await?; + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + let running_actions_manager = + Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { + root_action_directory: root_action_directory.clone(), + execution_configuration: ExecutionConfiguration::default(), + cas_store: cas_store.clone(), + ac_store: None, + historical_store: Store::new(cas_store.clone()), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::MAX, + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), + max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + active_input_leases: false, + directory_cache: None, + #[cfg(target_os = "linux")] + use_namespaces: use_namespaces(), + })?); + // The sleep inherits stdout and would hold it for a minute. + let command = Command { + arguments: vec![ + "sh".to_string(), + "-c".to_string(), + "echo partial; sleep 60 & exit 0".to_string(), + ], + working_directory: ".".to_string(), + environment_variables: vec![EnvironmentVariable { + name: "PATH".to_string(), + value: env::var("PATH").unwrap(), + }], + ..Default::default() + }; + let command_digest = serialize_and_upload_message( + &command, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let input_root_digest = serialize_and_upload_message( + &Directory::default(), + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let action = Action { + command_digest: Some(command_digest.into()), + input_root_digest: Some(input_root_digest.into()), + timeout: Some(prost_types::Duration { + seconds: 60, + nanos: 0, + }), + ..Default::default() + }; + let action_digest = serialize_and_upload_message( + &action, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let running_action_impl = running_actions_manager + .create_and_add_action( + WORKER_ID.to_string(), + StartExecute { + request_metadata: None, + execute_request: Some(ExecuteRequest { + action_digest: Some(action_digest.into()), + digest_function: ProtoDigestFunction::Sha256.into(), + ..Default::default() + }), + operation_id: OperationId::default().to_string(), + queued_timestamp: None, + platform: None, + worker_id: WORKER_ID.to_string(), + }, + ) + .await?; + let started = std::time::Instant::now(); + let result = run_action(running_action_impl).await?; + assert!( + started.elapsed() < Duration::from_secs(30), + "the action waited on the straggler: {:?}", + started.elapsed() + ); + assert_eq!(result.exit_code, 0, "{result:?}"); + assert!(result.error.is_none(), "{:?}", result.error); + let stdout = slow_store + .as_ref() + .get_part_unchunked(result.stdout_digest, 0, None) + .await?; + assert_eq!(stdout, "partial\n", "what it printed is kept"); + Ok(()) + } + + /// Each action gets its own `TMPDIR` under its action directory. + #[cfg(target_family = "unix")] + #[nativelink_test] + async fn set_tmpdir_gives_each_action_its_own_directory() + -> Result<(), Box> { + const WORKER_ID: &str = "foo_worker_id"; + let (_, slow_store, cas_store, _ac_store) = setup_stores().await?; + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + let running_actions_manager = + Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { + root_action_directory: root_action_directory.clone(), + execution_configuration: ExecutionConfiguration { + persistent_workers: PersistentWorkersSettings::default(), + set_tmpdir: true, + ..Default::default() + }, + cas_store: cas_store.clone(), + ac_store: None, + historical_store: Store::new(cas_store.clone()), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::MAX, + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), + max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + active_input_leases: false, + directory_cache: None, + #[cfg(target_os = "linux")] + use_namespaces: use_namespaces(), + })?); + let command = Command { + arguments: vec![ + "sh".to_string(), + "-c".to_string(), + "touch \"$TMPDIR/probe\" && printf %s \"$TMPDIR\" > out.txt".to_string(), + ], + working_directory: ".".to_string(), + environment_variables: vec![EnvironmentVariable { + name: "PATH".to_string(), + value: env::var("PATH").unwrap(), + }], + output_paths: vec!["out.txt".to_string()], + ..Default::default() + }; + let command_digest = serialize_and_upload_message( + &command, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let input_root_digest = serialize_and_upload_message( + &Directory::default(), + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let action = Action { + command_digest: Some(command_digest.into()), + input_root_digest: Some(input_root_digest.into()), + ..Default::default() + }; + let action_digest = serialize_and_upload_message( + &action, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let running_action_impl = running_actions_manager + .create_and_add_action( + WORKER_ID.to_string(), + StartExecute { + request_metadata: None, + execute_request: Some(ExecuteRequest { + action_digest: Some(action_digest.into()), + digest_function: ProtoDigestFunction::Sha256.into(), + ..Default::default() + }), + operation_id: OperationId::default().to_string(), + queued_timestamp: None, + platform: None, + worker_id: WORKER_ID.to_string(), + }, + ) + .await?; + let result = run_action(running_action_impl.clone()).await?; + assert_eq!(result.exit_code, 0, "{result:?}"); + let tmpdir = slow_store + .as_ref() + .get_part_unchunked(result.output_files[0].digest, 0, None) + .await?; + let tmpdir = String::from_utf8(tmpdir.to_vec())?; + assert!( + tmpdir.starts_with(&root_action_directory) && tmpdir.ends_with("/tmp"), + "TMPDIR should sit under the action directory, got {tmpdir}" + ); + let usage = running_action_impl + .resource_usage() + .expect("usage is reported for every executed action"); + assert_eq!( + usage.outcome, + ResourceOutcome::Completed as i32, + "{usage:?}" + ); + Ok(()) + } + + /// The outcome names who ended the action: the worker's own kills by + /// reason, a SIGKILL near the memory limit as the kernel, anything else + /// from outside. + #[nativelink_test] + async fn outcome_classification() -> Result<(), Box> { + const SIGKILL: i32 = 9; + const SIGTERM: i32 = 15; + assert_eq!( + classify_outcome(None, None, 900, Some(1000)), + (ResourceOutcome::Completed, false) + ); + assert_eq!( + classify_outcome(Some(KillReason::Memory), None, 0, None), + (ResourceOutcome::KilledMemory, true) + ); + assert_eq!( + classify_outcome(Some(KillReason::Timeout), Some(SIGKILL), 999, Some(1000)), + (ResourceOutcome::KilledTimeout, false) + ); + assert_eq!( + classify_outcome(Some(KillReason::External), Some(SIGTERM), 0, None), + (ResourceOutcome::KilledExternal, false) + ); + // The kernel's SIGKILL within 10% of the limit is an OOM kill. + assert_eq!( + classify_outcome(None, Some(SIGKILL), 900, Some(1000)), + (ResourceOutcome::KilledMemory, false) + ); + assert_eq!( + classify_outcome(None, Some(SIGKILL), 899, Some(1000)), + (ResourceOutcome::KilledExternal, false) + ); + // No yardstick, no verdict. + assert_eq!( + classify_outcome(None, Some(SIGKILL), 900, None), + (ResourceOutcome::KilledExternal, false) + ); + assert_eq!( + classify_outcome(None, Some(SIGTERM), 1000, Some(1000)), + (ResourceOutcome::KilledExternal, false) + ); + Ok(()) + } + + /// Output past the cap spills to disk and is uploaded whole; the worker + /// keeps only the cap in memory. + #[cfg(target_family = "unix")] + #[nativelink_test] + async fn output_past_the_cap_spills_to_disk_and_uploads_whole() + -> Result<(), Box> { + const WORKER_ID: &str = "foo_worker_id"; + const OUTPUT_BYTES: usize = 3_000_000; + let (_, slow_store, cas_store, _ac_store) = setup_stores().await?; + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + + let running_actions_manager = + Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { + root_action_directory: root_action_directory.clone(), + execution_configuration: ExecutionConfiguration { + persistent_workers: PersistentWorkersSettings::default(), + max_captured_output_bytes: 1_000_000, + ..Default::default() + }, + cas_store: cas_store.clone(), + ac_store: None, + historical_store: Store::new(cas_store.clone()), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::MAX, + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), + max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + active_input_leases: false, + directory_cache: None, + #[cfg(target_os = "linux")] + use_namespaces: use_namespaces(), + })?); + let command = Command { + arguments: vec![ + "sh".to_string(), + "-c".to_string(), + format!( + "head -c {OUTPUT_BYTES} /dev/zero | tr '\\0' a; head -c 100 /dev/zero | tr '\\0' b >&2" + ), + ], + working_directory: ".".to_string(), + environment_variables: vec![EnvironmentVariable { + name: "PATH".to_string(), + value: env::var("PATH").unwrap(), + }], + ..Default::default() + }; + let command_digest = serialize_and_upload_message( + &command, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let input_root_digest = serialize_and_upload_message( + &Directory::default(), + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let action = Action { + command_digest: Some(command_digest.into()), + input_root_digest: Some(input_root_digest.into()), + ..Default::default() + }; + let action_digest = serialize_and_upload_message( + &action, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let running_action_impl = running_actions_manager + .create_and_add_action( + WORKER_ID.to_string(), + StartExecute { + request_metadata: None, + execute_request: Some(ExecuteRequest { + action_digest: Some(action_digest.into()), + digest_function: ProtoDigestFunction::Sha256.into(), + ..Default::default() + }), + operation_id: OperationId::default().to_string(), + queued_timestamp: None, + platform: None, + worker_id: WORKER_ID.to_string(), + }, + ) + .await?; + let result = run_action(running_action_impl).await?; + assert_eq!(result.exit_code, 0); + assert_eq!(result.stdout_digest.size_bytes(), OUTPUT_BYTES as u64); + let stdout = slow_store + .as_ref() + .get_part_unchunked(result.stdout_digest, 0, None) + .await?; + assert_eq!(stdout.len(), OUTPUT_BYTES); + assert!( + stdout.iter().all(|b| *b == b'a'), + "spilled output should be intact" + ); + assert_eq!( + result.stderr_digest.size_bytes(), + 100, + "small stderr stays in memory" + ); + Ok(()) + } + + /// The log excerpt of a failed command is cut by bytes: through a + /// multibyte character it ends in U+FFFD instead of panicking, and + /// bytes that were never text are shown the same way. + #[nativelink_test] + async fn log_excerpt_cuts_multibyte_output_safely() -> Result<(), Box> { + let text = "héllo wörld"; + assert_eq!(log_excerpt(text.as_bytes(), 2), "h\u{FFFD}"); + assert_eq!(log_excerpt(text.as_bytes(), 3), "hé"); + assert_eq!(log_excerpt(text.as_bytes(), 1000), text); + assert_eq!(log_excerpt(&[0xff, b'a'], 1000), "\u{FFFD}a"); + assert_eq!(log_excerpt(b"", 1000), ""); + Ok(()) + } + + /// A directory under the action root that no running action owns and + /// The sweep visits a directory whose operation is already marked as + /// cleaning (its own cleanup is running). Taking the mark a second time + /// must hand back `None`, not a guard built under the lock and dropped + /// there, which deadlocked the runtime thread on its own mutex. + #[nativelink_test] + async fn orphan_sweep_skips_a_directory_already_being_cleaned_without_deadlocking() + -> Result<(), Box> { + let (_, _, cas_store, _ac_store) = setup_stores().await?; + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + let running_actions_manager = + Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { + root_action_directory: root_action_directory.clone(), + execution_configuration: ExecutionConfiguration::default(), + cas_store: cas_store.clone(), + ac_store: None, + historical_store: Store::new(cas_store.clone()), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::MAX, + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), + max_cleanup_wait: Duration::ZERO, + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + active_input_leases: false, + directory_cache: None, + #[cfg(target_os = "linux")] + use_namespaces: use_namespaces(), + })?); + let operation_id = OperationId::from("mid-cleanup"); + fs::create_dir_all(format!("{root_action_directory}/{operation_id}")).await?; + let held = running_actions_manager + .perform_cleanup(operation_id.clone()) + .expect("first mark is granted"); + // The second mark and the sweep on their own thread: a deadlock + // blocks that thread, not this one, and the wait turns it into a + // failure instead of a hung test. + let (tx, rx) = std::sync::mpsc::channel(); + let manager = running_actions_manager.clone(); + let marked_operation_id = operation_id.clone(); + std::thread::spawn(move || { + let second_mark = manager.perform_cleanup(marked_operation_id).is_some(); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("runtime"); + let swept = runtime.block_on(manager.sweep_orphaned_directories()); + drop(tx.send((second_mark, swept.map_err(|e| e.to_string())))); + }); + let (second_mark, swept) = rx + .recv_timeout(Duration::from_secs(10)) + .expect("the second cleanup mark deadlocked on its own lock"); + assert!(!second_mark, "a mark already held is not granted again"); + assert_eq!( + swept, + Ok(0), + "a directory being cleaned is left to its cleanup" + ); + assert!( + Path::new(&format!("{root_action_directory}/{operation_id}")).exists(), + "the sweep must not remove a directory being cleaned" + ); + drop(held); + assert_eq!( + running_actions_manager.sweep_orphaned_directories().await?, + 1, + "once the mark is released the settled directory is an orphan" + ); + Ok(()) + } + + /// A manager with a root directory, a cleanup wait of two seconds and no + /// directory cache. + async fn manager_for_cleanup_tests( + root_action_directory: &str, + ) -> Result<(Arc, Arc), Box> + { + let (_, slow_store, cas_store, _ac_store) = setup_stores().await?; + let running_actions_manager = + Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { + root_action_directory: root_action_directory.to_string(), + execution_configuration: ExecutionConfiguration::default(), + cas_store: cas_store.clone(), + ac_store: None, + historical_store: Store::new(cas_store), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::MAX, + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), + max_cleanup_wait: Duration::from_secs(2), + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + active_input_leases: false, + directory_cache: None, + #[cfg(target_os = "linux")] + use_namespaces: use_namespaces(), + })?); + Ok((running_actions_manager, slow_store)) + } + + async fn empty_action_digest( + store: &Arc, + ) -> Result> { + let command_digest = serialize_and_upload_message( + &Command::default(), + store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let input_root_digest = serialize_and_upload_message( + &Directory::default(), + store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + Ok(serialize_and_upload_message( + &Action { + command_digest: Some(command_digest.into()), + input_root_digest: Some(input_root_digest.into()), + ..Default::default() + }, + store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?) + } + + fn start_execute_with(action_digest: DigestInfo, operation_id: &OperationId) -> StartExecute { + StartExecute { + request_metadata: None, + execute_request: Some(ExecuteRequest { + action_digest: Some(action_digest.into()), + digest_function: ProtoDigestFunction::Sha256.into(), + ..Default::default() + }), + operation_id: operation_id.to_string(), + ..Default::default() + } + } + + /// An action's cleanup waits for a held mark (the sweep's, a retry's) + /// and then cleans; it does not step aside and leave its entries behind. + #[nativelink_test] + async fn cleanup_waits_for_a_held_mark_then_cleans() -> Result<(), Box> + { + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + let (running_actions_manager, store) = + manager_for_cleanup_tests(&root_action_directory).await?; + let action_digest = empty_action_digest(&store).await?; + let operation_id = OperationId::default(); + let action = running_actions_manager + .create_and_add_action( + "test-worker".to_string(), + start_execute_with(action_digest, &operation_id), + ) + .await?; + let action_directory = format!("{root_action_directory}/{operation_id}"); + assert!(Path::new(&action_directory).exists()); + let held = running_actions_manager + .perform_cleanup(operation_id.clone()) + .expect("the mark is free"); + let mut cleanup = tokio::spawn(action.cleanup()); + tokio::select! { + _ = &mut cleanup => panic!("cleanup ran while the mark was held"), + () = tokio::time::sleep(Duration::from_millis(300)) => {} + } + assert!( + Path::new(&action_directory).exists(), + "nothing is removed while the mark is held" + ); + drop(held); + tokio::time::timeout(Duration::from_secs(5), cleanup) + .await + .expect("cleanup did not resume when the mark was released")??; + assert!( + !Path::new(&action_directory).exists(), + "the directory goes once cleanup runs" + ); + // The manager entry went with it: nothing for kill_all to wait on. + tokio::time::timeout(Duration::from_secs(5), running_actions_manager.kill_all()) + .await + .expect("an entry was left behind"); + Ok(()) + } + + /// A retry removes the stale directory of its earlier attempt only with + /// the mark held, so it waits while the sweep (or anyone) holds it and + /// never races a removal of the same tree. + #[nativelink_test] + async fn retry_waits_for_the_mark_before_removing_a_stale_directory() + -> Result<(), Box> { + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + let (running_actions_manager, store) = + manager_for_cleanup_tests(&root_action_directory).await?; + let action_digest = empty_action_digest(&store).await?; + let operation_id = OperationId::default(); + let stale = format!("{root_action_directory}/{operation_id}"); + fs::create_dir_all(&stale).await?; + tokio::fs::write(format!("{stale}/left-behind"), b"x").await?; + let held = running_actions_manager + .perform_cleanup(operation_id.clone()) + .expect("the mark is free"); + let manager = running_actions_manager.clone(); + let start_execute = start_execute_with(action_digest, &operation_id); + let mut retry = tokio::spawn(async move { + manager + .create_and_add_action("w".to_string(), start_execute) + .await + }); + tokio::select! { + _ = &mut retry => panic!("the retry proceeded while the mark was held"), + () = tokio::time::sleep(Duration::from_millis(300)) => {} + } + assert!( + Path::new(&format!("{stale}/left-behind")).exists(), + "the stale tree is untouched while the mark is held" + ); + drop(held); + let action = tokio::time::timeout(Duration::from_secs(5), retry) + .await + .expect("the retry did not proceed once the mark was released")??; + assert!( + !Path::new(&format!("{stale}/left-behind")).exists(), + "the stale tree was replaced by the retry's fresh directory" + ); + action.cleanup().await?; + Ok(()) + } + + /// that has settled is removed by the sweep; a running action's is kept. + #[nativelink_test] + async fn orphan_sweep_removes_settled_unowned_directories() + -> Result<(), Box> { + let (_, _, cas_store, _ac_store) = setup_stores().await?; + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + let running_actions_manager = + Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { + root_action_directory: root_action_directory.clone(), + execution_configuration: ExecutionConfiguration::default(), + cas_store: cas_store.clone(), + ac_store: None, + historical_store: Store::new(cas_store.clone()), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::MAX, + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), + // Zero: anything unowned counts as settled at once. + max_cleanup_wait: Duration::ZERO, + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + active_input_leases: false, + directory_cache: None, + #[cfg(target_os = "linux")] + use_namespaces: use_namespaces(), + })?); + let stray = format!("{root_action_directory}/left-behind-by-a-failed-cleanup"); + fs::create_dir_all(&stray).await?; + let stray_file = format!("{root_action_directory}/not-a-directory"); + tokio::fs::write(&stray_file, b"x").await?; + + assert_eq!( + running_actions_manager.sweep_orphaned_directories().await?, + 1 + ); + assert!(!Path::new(&stray).exists(), "the orphan should be gone"); + assert!( + Path::new(&stray_file).exists(), + "files are not the sweep's business" + ); + assert_eq!( + running_actions_manager.sweep_orphaned_directories().await?, + 0 + ); + Ok(()) + } + + /// A symlink under the action root is not an action directory, whatever + /// it points at: the sweep neither follows it nor removes it. + #[cfg(target_family = "unix")] + #[nativelink_test] + async fn orphan_sweep_leaves_symlinks_alone() -> Result<(), Box> { + let (_, _, cas_store, _ac_store) = setup_stores().await?; + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + let running_actions_manager = + Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { + root_action_directory: root_action_directory.clone(), + execution_configuration: ExecutionConfiguration::default(), + cas_store: cas_store.clone(), + ac_store: None, + historical_store: Store::new(cas_store.clone()), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::MAX, + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), + max_cleanup_wait: Duration::ZERO, + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + active_input_leases: false, + directory_cache: None, + #[cfg(target_os = "linux")] + use_namespaces: use_namespaces(), + })?); + // A settled directory outside the root, reachable through a symlink + // inside it, next to a real orphan. + let elsewhere = make_temp_path("elsewhere"); + fs::create_dir_all(&elsewhere).await?; + tokio::fs::write(format!("{elsewhere}/kept"), b"x").await?; + let link = format!("{root_action_directory}/a-link"); + tokio::fs::symlink(&elsewhere, &link).await?; + let orphan = format!("{root_action_directory}/b-orphan"); + fs::create_dir_all(&orphan).await?; + + assert_eq!( + running_actions_manager.sweep_orphaned_directories().await?, + 1 + ); + assert!(!Path::new(&orphan).exists(), "the orphan should be gone"); + assert!( + Path::new(&link).symlink_metadata().is_ok(), + "the symlink is not the sweep's business" + ); + assert!( + Path::new(&format!("{elsewhere}/kept")).exists(), + "and neither is what it points at" + ); + Ok(()) + } + + #[cfg(target_os = "linux")] + #[nativelink_test] + async fn action_over_its_memory_reservation_is_killed_by_the_worker() + -> Result<(), Box> { + const WORKER_ID: &str = "foo_worker_id"; + let (_, _, cas_store, _ac_store) = setup_stores().await?; + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + + let running_actions_manager = + Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { + root_action_directory: root_action_directory.clone(), + execution_configuration: ExecutionConfiguration { + resource_enforcement: Some(ResourceEnforcement { + memory: Some(MemoryLimit { + property_name: "memory_kb".to_string(), + headroom_percent: 20, + }), + disk_property_name: None, + disk_soft: false, + disk_headroom_percent: 20, + }), + ..Default::default() + }, + cas_store: cas_store.clone(), + ac_store: None, + historical_store: Store::new(cas_store.clone()), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::MAX, + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), + max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + active_input_leases: false, + directory_cache: None, + #[cfg(target_os = "linux")] + use_namespaces: use_namespaces(), + })?); + // A shell holding 60 MB in a variable, then a child that would keep + // the memory if only the shell were killed. + let command = Command { + arguments: vec![ + "sh".to_string(), + "-c".to_string(), + "s=$(head -c 60000000 /dev/zero | tr '\\0' a); sleep 30; echo $s | head -c 1" + .to_string(), + ], + working_directory: ".".to_string(), + environment_variables: vec![EnvironmentVariable { + name: "PATH".to_string(), + value: env::var("PATH").unwrap(), + }], + ..Default::default() + }; + let command_digest = serialize_and_upload_message( + &command, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let input_root_digest = serialize_and_upload_message( + &Directory::default(), + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + // The client tagged nothing; the reservation arrives the way a hint or + // cold start does, on the scheduler's StartExecute, not on the Action. + let action = Action { + command_digest: Some(command_digest.into()), + input_root_digest: Some(input_root_digest.into()), + ..Default::default() + }; + let action_digest = serialize_and_upload_message( + &action, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + + let running_action_impl = running_actions_manager + .create_and_add_action( + WORKER_ID.to_string(), + StartExecute { + request_metadata: None, + execute_request: Some(ExecuteRequest { + action_digest: Some(action_digest.into()), + digest_function: ProtoDigestFunction::Sha256.into(), + ..Default::default() + }), + operation_id: OperationId::default().to_string(), + queued_timestamp: None, + // 2 MiB reserved: the shell is far over it within two samples. + platform: Some(Platform { + properties: vec![Property { + name: "memory_kb".into(), + value: "2048".into(), + }], + }), + worker_id: WORKER_ID.to_string(), + }, + ) + .await?; + + let started = std::time::Instant::now(); + let result = run_action(running_action_impl.clone()).await?; + let err = result + .error + .expect("an action over its reservation must fail with an error"); + assert_eq!(err.code, Code::FailedPrecondition, "{err}"); + assert!( + err.to_string().contains("memory reservation"), + "error should name the reservation: {err}" + ); + assert!( + started.elapsed() < Duration::from_secs(20), + "the worker should have killed the action long before its 30s sleep ended" + ); + // The measurement says what happened and what it was measured against. + let usage = running_action_impl + .resource_usage() + .expect("usage is reported for every executed action"); + assert_eq!( + usage.outcome, + ResourceOutcome::KilledMemory as i32, + "{usage:?}" + ); + assert!(usage.enforced, "the worker's own limit ended it: {usage:?}"); + assert_eq!( + usage.reserved.map(|r| r.memory_kb), + Some(2048), + "the reservation travels with the measurement: {usage:?}" + ); + // No wall time assertion: the gross-overrun kill lands within a millisecond. + Ok(()) + } + + #[cfg_attr(feature = "nix", ignore)] + #[nativelink_test] + async fn entrypoint_sends_timeout_via_side_channel() -> Result<(), Box> + { + #[cfg(target_family = "unix")] + const TEST_WRAPPER_SCRIPT_CONTENT: &str = "\ +#!/bin/bash +echo '{\"failure\":\"timeout\"}' > \"$SIDE_CHANNEL_FILE\" +exit 1 +"; + #[cfg(target_family = "windows")] + const TEST_WRAPPER_SCRIPT_CONTENT: &str = "\ +@echo off +echo | set /p={\"failure\":\"timeout\"} 1>&2 > %SIDE_CHANNEL_FILE% +exit 1 +"; + const WORKER_ID: &str = "foo_worker_id"; + + let (_, _, cas_store, ac_store) = setup_stores().await?; + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + + let test_wrapper_script = { + let test_wrapper_dir = make_temp_path("wrapper_dir"); + fs::create_dir_all(&test_wrapper_dir).await?; + #[cfg(target_family = "unix")] + let test_wrapper_script = OsString::from(test_wrapper_dir + "/test_wrapper_script.sh"); + #[cfg(target_family = "windows")] + let test_wrapper_script = + OsString::from(test_wrapper_dir + "\\test_wrapper_script.bat"); + { + let mut file_options = std::fs::OpenOptions::new(); + file_options.create(true); + file_options.truncate(true); + file_options.write(true); + #[cfg(target_family = "unix")] + file_options.mode(0o777); + let mut test_wrapper_script_handle = file_options + .open(OsString::from(&test_wrapper_script)) + .unwrap(); + test_wrapper_script_handle + .write_all(TEST_WRAPPER_SCRIPT_CONTENT.as_bytes()) + .unwrap(); + test_wrapper_script_handle.sync_all().unwrap(); + // Note: Github runners appear to use some kind of filesystem driver + // that does not sync data as expected. This is the easiest solution. + // See: https://github.com/pantsbuild/pants/issues/10507 + // See: https://github.com/moby/moby/issues/9547 + std::process::Command::new("sync").output().unwrap(); + } + test_wrapper_script + }; + + let running_actions_manager = + Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { + root_action_directory: root_action_directory.clone(), + execution_configuration: ExecutionConfiguration { + persistent_workers: PersistentWorkersSettings::default(), + max_captured_output_bytes: 0, + resource_enforcement: None, + kill_grace: Duration::ZERO, + set_tmpdir: false, + buck2_file_capture: None, + entrypoint: Some(test_wrapper_script.into_string().unwrap()), + additional_environment: Some(HashMap::from([( + "SIDE_CHANNEL_FILE".to_string(), + EnvironmentSource::SideChannelFile, + )])), + }, + cas_store: cas_store.clone(), + ac_store: Some(Store::new(ac_store.clone())), + historical_store: Store::new(cas_store.clone()), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::MAX, + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), + max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + active_input_leases: false, + directory_cache: None, + #[cfg(target_os = "linux")] + use_namespaces: use_namespaces(), + })?); + let arguments = vec!["true".to_string()]; + let command = Command { + arguments, + working_directory: ".".to_string(), + environment_variables: vec![EnvironmentVariable { + name: "PATH".to_string(), + value: env::var("PATH").unwrap(), + }], + ..Default::default() + }; + let command_digest = serialize_and_upload_message( + &command, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let input_root_digest = serialize_and_upload_message( + &Directory::default(), + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let action = Action { + command_digest: Some(command_digest.into()), + input_root_digest: Some(input_root_digest.into()), + ..Default::default() + }; + let action_digest = serialize_and_upload_message( + &action, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + + let execute_request = ExecuteRequest { + action_digest: Some(action_digest.into()), + digest_function: ProtoDigestFunction::Sha256.into(), + ..Default::default() + }; + let operation_id = OperationId::default().to_string(); + + let running_action_impl = running_actions_manager + .clone() + .create_and_add_action( + WORKER_ID.to_string(), + StartExecute { + request_metadata: None, + execute_request: Some(execute_request), + operation_id, + queued_timestamp: Some(make_system_time(1000).into()), + platform: action.platform.clone(), + worker_id: WORKER_ID.to_string(), + }, + ) + .await?; + + let result = run_action(running_action_impl).await?; + assert_eq!(result.exit_code, 1, "Exit code should be 1"); + assert_eq!( + result.error.err_tip(|| "Error should exist")?.code, + Code::DeadlineExceeded + ); + assert!(logs_contain( + "Command returned non-zero exit code exit_code=1 stdout=\"\" stderr=\"\" command=[\"true\"]" + )); + Ok(()) + } + + #[nativelink_test] + async fn caches_results_in_action_cache_store() -> Result<(), Box> { + let (_, _, cas_store, ac_store) = setup_stores().await?; let running_actions_manager = Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { @@ -3092,6 +5509,7 @@ exit 1 }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -3171,6 +5589,7 @@ exit 1 }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -3257,6 +5676,7 @@ exit 1 }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -3364,6 +5784,7 @@ exit 1 }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -3415,6 +5836,7 @@ exit 1 }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -3487,6 +5909,7 @@ exit 1 }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -3610,6 +6033,7 @@ exit 1 }, max_action_timeout: MAX_TIMEOUT_DURATION, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -3703,6 +6127,7 @@ exit 1 }, max_action_timeout: MAX_TIMEOUT_DURATION, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -3762,7 +6187,7 @@ exit 1 ); } { - // Ensure we reject tasks that have a timeout set too high. + // A timeout over the maximum is clamped to it, not rejected. static SENT_TIMEOUT: AtomicI64 = AtomicI64::new(-1); const MAX_TIMEOUT_DURATION: Duration = Duration::from_secs(100); const TASK_TIMEOUT: Duration = Duration::from_secs(200); @@ -3796,6 +6221,7 @@ exit 1 }, max_action_timeout: MAX_TIMEOUT_DURATION, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -3848,8 +6274,12 @@ exit 1 }) }) .await; - assert_eq!(SENT_TIMEOUT.load(Ordering::Relaxed), -1); - assert_eq!(result.err().unwrap().code, Code::InvalidArgument); + assert!(result.is_ok(), "{result:?}"); + assert_eq!( + SENT_TIMEOUT.load(Ordering::Relaxed), + i64::try_from(MAX_TIMEOUT_DURATION.as_millis()) + .expect("MAX_TIMEOUT_DURATION.as_millis() exceeds i64::MAX") + ); } Ok(()) } @@ -3886,6 +6316,7 @@ exit 1 }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -4047,6 +6478,7 @@ exit 1 }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -4173,17 +6605,17 @@ exit 1 let result = execute_results_fut.await; kill_all_fut.await; { - // Ensure our results are correct. - let action_result = result?; - let err = action_result - .error - .as_ref() - .err_tip(|| format!("No error exists in result : {action_result:?}"))?; - assert_eq!( - err.code, - Code::Aborted, - "Expected Aborted : {action_result:?}" - ); + // Ensure our results are correct. The kill lands on the command + // or, when the fetch has not finished yet, on the fetch; either + // way the action ends as aborted. + let err = match result { + Ok(action_result) => action_result + .error + .clone() + .err_tip(|| format!("No error exists in result : {action_result:?}"))?, + Err(err) => err, + }; + assert_eq!(err.code, Code::Aborted, "Expected Aborted : {err:?}"); } Ok(()) @@ -4223,6 +6655,7 @@ exit 1 }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -4336,6 +6769,7 @@ exit 1 }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -4451,6 +6885,7 @@ exit 1 }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -4544,6 +6979,136 @@ exit 1 Ok(()) } + #[cfg(target_os = "linux")] + #[nativelink_test] + async fn isolate_tmp_gives_action_a_private_tmp() -> Result<(), Box> { + const WORKER_ID: &str = "foo_worker_id"; + + if !namespace_utils::namespaces_supported(true, true) { + return Ok(()); + } + + let (_, _, cas_store, ac_store) = setup_stores().await?; + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + + let running_actions_manager = + Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { + root_action_directory, + cas_store: cas_store.clone(), + ac_store: Some(Store::new(ac_store.clone())), + execution_configuration: ExecutionConfiguration::default(), + historical_store: Store::new(cas_store.clone()), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::MAX, + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), + max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + active_input_leases: false, + directory_cache: None, + use_namespaces: + nativelink_worker::running_actions_manager::UseNamespaces::YesAndMount { + isolate_tmp: true, + }, + })?); + + // A file in the host's /tmp must be invisible to the action, and the + // action's own scratch file in /tmp must never reach the host, while + // outputs in the work directory still get uploaded. + let unique = format!( + "{}_{}", + std::process::id(), + SystemTime::now().duration_since(UNIX_EPOCH)?.as_nanos() + ); + let host_tmp_file = format!("/tmp/nativelink_test_host_{unique}"); + let scratch_file = format!("/tmp/nativelink_test_scratch_{unique}"); + tokio::fs::write(&host_tmp_file, "host").await?; + + let command = Command { + arguments: vec![ + "sh".to_string(), + "-c".to_string(), + format!( + "test ! -e {host_tmp_file} && echo scratch > {scratch_file} && test -f {scratch_file} && echo ok > out.txt" + ), + ], + output_paths: vec!["out.txt".to_string()], + working_directory: ".".to_string(), + environment_variables: vec![EnvironmentVariable { + name: "PATH".to_string(), + value: env::var("PATH").unwrap(), + }], + ..Default::default() + }; + let command_digest = serialize_and_upload_message( + &command, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let input_root_digest = serialize_and_upload_message( + &Directory::default(), + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let action = Action { + command_digest: Some(command_digest.into()), + input_root_digest: Some(input_root_digest.into()), + ..Default::default() + }; + let action_digest = serialize_and_upload_message( + &action, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + + let execute_request = ExecuteRequest { + action_digest: Some(action_digest.into()), + digest_function: ProtoDigestFunction::Sha256.into(), + ..Default::default() + }; + let operation_id = OperationId::default().to_string(); + let running_action_impl = running_actions_manager + .create_and_add_action( + WORKER_ID.to_string(), + StartExecute { + execute_request: Some(execute_request), + operation_id, + ..Default::default() + }, + ) + .await?; + + let action_result = run_action(running_action_impl.clone()).await; + tokio::fs::remove_file(&host_tmp_file).await?; + let action_result = action_result?; + assert_eq!( + action_result.exit_code, 0, + "action should neither see the host's /tmp nor fail to write its own" + ); + assert_eq!( + action_result.output_files.len(), + 1, + "out.txt should have been uploaded from the work directory" + ); + assert_eq!( + action_result.output_files[0].name_or_path, + NameOrPath::Path("out.txt".to_string()) + ); + assert!( + tokio::fs::metadata(&scratch_file).await.is_err(), + "the action's /tmp scratch file leaked to the host" + ); + Ok(()) + } + #[nativelink_test] async fn action_directory_contents_are_cleaned() -> Result<(), Box> { const WORKER_ID: &str = "foo_worker_id"; @@ -4567,6 +7132,7 @@ exit 1 }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -4656,7 +7222,7 @@ exit 1 use nativelink_config::cas_server::Buck2FileCaptureConfig; use nativelink_proto::build::bazel::remote::execution::v2::{RequestMetadata, ToolDetails}; - async fn wait_for_file(path: &std::path::Path) -> Result<(), tokio::time::error::Elapsed> { + async fn wait_for_file(path: &Path) -> Result<(), tokio::time::error::Elapsed> { tokio::time::timeout(Duration::from_secs(10), async { while !path.exists() { tokio::time::sleep(Duration::from_millis(5)).await; @@ -4708,6 +7274,8 @@ while [ ! -f "$0.d/release" ]; do "{sleep}" 0.01; done let manager = Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { root_action_directory: actions.to_string_lossy().into_owned(), execution_configuration: ExecutionConfiguration { + persistent_workers: PersistentWorkersSettings::default(), + resource_enforcement: None, buck2_file_capture: Some(Buck2FileCaptureConfig { executable: helper.to_string_lossy().into_owned(), gateway: "http://unused-test-gateway".to_string(), @@ -4729,6 +7297,7 @@ while [ ! -f "$0.d/release" ]; do "{sleep}" 0.01; done }, max_action_timeout: Duration::from_mins(1), max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -4860,6 +7429,12 @@ while [ ! -f "$0.d/release" ]; do "{sleep}" 0.01; done ) -> Result<(Action, DigestInfo), Error> { let command = Command { arguments: vec![format!("./{WORKER_SCRIPT_NAME}"), "@args.txt".to_string()], + // The worker process gets the action's environment and no + // more, as Bazel sends it; the script needs PATH for sed. + environment_variables: vec![EnvironmentVariable { + name: "PATH".to_string(), + value: env::var("PATH").unwrap_or_default(), + }], output_paths: vec!["count.txt".to_string()], platform: Some(Platform { properties: vec![ @@ -4974,6 +7549,7 @@ while [ ! -f "$0.d/release" ]; do "{sleep}" 0.01; done }, max_action_timeout: Duration::from_secs(30), max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -5074,6 +7650,7 @@ done }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -5260,6 +7837,7 @@ done }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -5374,6 +7952,11 @@ done Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { root_action_directory: root_action_directory.clone(), execution_configuration: ExecutionConfiguration { + persistent_workers: PersistentWorkersSettings::default(), + max_captured_output_bytes: 0, + resource_enforcement: None, + kill_grace: Duration::ZERO, + set_tmpdir: false, buck2_file_capture: None, entrypoint: None, additional_environment: None, @@ -5387,6 +7970,7 @@ done }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -5522,6 +8106,11 @@ done Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs { root_action_directory: root_action_directory.clone(), execution_configuration: ExecutionConfiguration { + persistent_workers: PersistentWorkersSettings::default(), + max_captured_output_bytes: 0, + resource_enforcement: None, + kill_grace: Duration::ZERO, + set_tmpdir: false, buck2_file_capture: None, entrypoint: None, additional_environment: None, @@ -5535,6 +8124,7 @@ done }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -5650,6 +8240,7 @@ done }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -5795,6 +8386,7 @@ done }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -5971,6 +8563,7 @@ done }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::MAX, + max_download_timeout: Duration::MAX, max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, @@ -6131,6 +8724,7 @@ done }, max_action_timeout: Duration::MAX, max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), timeout_handled_externally: false, diff --git a/nativelink-worker/tests/utils/local_worker_test_utils.rs b/nativelink-worker/tests/utils/local_worker_test_utils.rs index 332832b2532..e143a314628 100644 --- a/nativelink-worker/tests/utils/local_worker_test_utils.rs +++ b/nativelink-worker/tests/utils/local_worker_test_utils.rs @@ -22,14 +22,14 @@ use hyper::body::Frame; use nativelink_config::cas_server::{EndpointConfig, LocalWorkerConfig, WorkerProperty}; use nativelink_error::{Error, make_err}; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{ - ConnectWorkerRequest, ExecuteComplete, ExecuteResult, GoingAwayRequest, KeepAliveRequest, - UpdateForWorker, + ConnectWorkerRequest, ExecuteAccepted, ExecuteComplete, ExecuteDeclined, ExecuteResult, + GoingAwayRequest, KeepAliveRequest, UpdateForWorker, }; use nativelink_util::channel_body_for_tests::ChannelBody; use nativelink_util::shutdown_guard::ShutdownGuard; use nativelink_util::spawn; use nativelink_util::task::JoinHandleDropGuard; -use nativelink_worker::local_worker::LocalWorker; +use nativelink_worker::local_worker::{LocalWorker, WorkerRegistration}; use nativelink_worker::worker_api_client_wrapper::WorkerApiClientTrait; use tokio::sync::{broadcast, mpsc}; use tonic::{Code, Status}; @@ -56,6 +56,8 @@ const BROADCAST_CAPACITY: usize = 1; enum WorkerClientApiCalls { ConnectWorker(ConnectWorkerRequest), ExecutionResponse(ExecuteResult), + ExecuteAccepted(ExecuteAccepted), + ExecuteDeclined(ExecuteDeclined), } #[derive(Debug)] @@ -66,6 +68,8 @@ enum WorkerClientApiCalls { enum WorkerClientApiReturns { ConnectWorker(Result>, Status>), ExecutionResponse(Result<(), Error>), + ExecuteAccepted(Result<(), Error>), + ExecuteDeclined(Result<(), Error>), } #[derive(Clone)] @@ -76,6 +80,7 @@ pub(crate) struct MockWorkerApiClient { tx_resp: mpsc::UnboundedSender, keep_alives_count: u8, pub going_away_count: Arc, + pub going_away_drain: Arc, pub execution_complete_count: Arc, } @@ -90,6 +95,7 @@ impl MockWorkerApiClient { tx_resp, keep_alives_count: 0, going_away_count: Arc::new(AtomicU64::new(0)), + going_away_drain: Arc::new(std::sync::atomic::AtomicBool::new(false)), execution_complete_count: Arc::new(AtomicU64::new(0)), } } @@ -113,7 +119,7 @@ impl MockWorkerApiClient { .expect("Could not receive msg in mpsc") { WorkerClientApiCalls::ConnectWorker(req) => req, - req @ WorkerClientApiCalls::ExecutionResponse(_) => { + req => { panic!("expect_connect_worker expected ConnectWorker, got : {req:?}") } }; @@ -123,6 +129,44 @@ impl MockWorkerApiClient { req } + pub(crate) async fn expect_execute_accepted( + &self, + result: Result<(), Error>, + ) -> ExecuteAccepted { + let mut rx_call_lock = self.rx_call.lock().await; + let req = match rx_call_lock + .recv() + .await + .expect("Could not receive msg in mpsc") + { + WorkerClientApiCalls::ExecuteAccepted(req) => req, + req => panic!("expect_execute_accepted expected ExecuteAccepted, got : {req:?}"), + }; + self.tx_resp + .send(WorkerClientApiReturns::ExecuteAccepted(result)) + .expect("Could not send request to mpsc"); + req + } + + pub(crate) async fn expect_execute_declined( + &self, + result: Result<(), Error>, + ) -> ExecuteDeclined { + let mut rx_call_lock = self.rx_call.lock().await; + let req = match rx_call_lock + .recv() + .await + .expect("Could not receive msg in mpsc") + { + WorkerClientApiCalls::ExecuteDeclined(req) => req, + req => panic!("expect_execute_declined expected ExecuteDeclined, got : {req:?}"), + }; + self.tx_resp + .send(WorkerClientApiReturns::ExecuteDeclined(result)) + .expect("Could not send request to mpsc"); + req + } + pub(crate) async fn expect_execution_response( &self, result: Result<(), Error>, @@ -134,7 +178,7 @@ impl MockWorkerApiClient { .expect("Could not receive msg in mpsc") { WorkerClientApiCalls::ExecutionResponse(req) => req, - req @ WorkerClientApiCalls::ConnectWorker(_) => { + req => { panic!("expect_execution_response expected ExecutionResponse, got : {req:?}") } }; @@ -160,7 +204,7 @@ impl WorkerApiClientTrait for MockWorkerApiClient { .expect("Could not receive msg in mpsc") { WorkerClientApiReturns::ConnectWorker(result) => result, - resp @ WorkerClientApiReturns::ExecutionResponse(_) => { + resp => { panic!("connect_worker expected ConnectWorker response, received {resp:?}") } } @@ -176,8 +220,10 @@ impl WorkerApiClientTrait for MockWorkerApiClient { } } - async fn going_away(&mut self, _request: GoingAwayRequest) -> Result<(), Error> { + async fn going_away(&mut self, request: GoingAwayRequest) -> Result<(), Error> { self.going_away_count.fetch_add(1, Ordering::Relaxed); + self.going_away_drain + .store(request.drain, Ordering::Relaxed); Ok(()) } @@ -192,7 +238,7 @@ impl WorkerApiClientTrait for MockWorkerApiClient { .expect("Could not receive msg in mpsc") { WorkerClientApiReturns::ExecutionResponse(result) => result, - resp @ WorkerClientApiReturns::ConnectWorker(_) => { + resp => { panic!("execution_response expected ExecutionResponse response, received {resp:?}") } } @@ -203,6 +249,36 @@ impl WorkerApiClientTrait for MockWorkerApiClient { .fetch_add(1, Ordering::Relaxed); Ok(()) } + + async fn execute_accepted(&mut self, request: ExecuteAccepted) -> Result<(), Error> { + self.tx_call + .send(WorkerClientApiCalls::ExecuteAccepted(request)) + .expect("Could not send request to mpsc"); + let mut rx_resp_lock = self.rx_resp.lock().await; + match rx_resp_lock + .recv() + .await + .expect("Could not receive msg in mpsc") + { + WorkerClientApiReturns::ExecuteAccepted(result) => result, + resp => panic!("execute_accepted expected ExecuteAccepted response, received {resp:?}"), + } + } + + async fn execute_declined(&mut self, request: ExecuteDeclined) -> Result<(), Error> { + self.tx_call + .send(WorkerClientApiCalls::ExecuteDeclined(request)) + .expect("Could not send request to mpsc"); + let mut rx_resp_lock = self.rx_resp.lock().await; + match rx_resp_lock + .recv() + .await + .expect("Could not receive msg in mpsc") + { + WorkerClientApiReturns::ExecuteDeclined(result) => result, + resp => panic!("execute_declined expected ExecuteDeclined response, received {resp:?}"), + } + } } pub(crate) fn setup_grpc_stream() -> ( @@ -236,6 +312,7 @@ pub(crate) async fn setup_local_worker_with_config( ); let (shutdown_tx_test, _) = broadcast::channel::(BROADCAST_CAPACITY); + let registration = worker.registration(); let drop_guard = spawn!("local_worker_spawn", async move { worker.run(shutdown_tx_test.subscribe()).await }); @@ -247,6 +324,7 @@ pub(crate) async fn setup_local_worker_with_config( maybe_streaming_response: Some(streaming_response), maybe_tx_stream: Some(tx_stream), + registration, drop_guard, } @@ -273,6 +351,8 @@ pub(crate) struct TestContext { pub maybe_streaming_response: Option>>, pub maybe_tx_stream: Option>>, + /// The flag the worker's readiness reads. + pub registration: Arc, drop_guard: JoinHandleDropGuard>, } diff --git a/src/bin/nativelink.rs b/src/bin/nativelink.rs index d354a0f9cec..55022b1ee1f 100644 --- a/src/bin/nativelink.rs +++ b/src/bin/nativelink.rs @@ -44,7 +44,7 @@ use nativelink_service::capabilities_server::CapabilitiesServer; use nativelink_service::cas_server::CasServer; use nativelink_service::execution_server::ExecutionServer; use nativelink_service::fetch_server::FetchServer; -use nativelink_service::health_server::HealthServer; +use nativelink_service::health_server::{HealthServer, health_paths}; use nativelink_service::push_server::PushServer; use nativelink_service::wire_compression::RemoteCacheCompressionInstances; use nativelink_service::worker_api_server::WorkerApiServer; @@ -63,7 +63,7 @@ use nativelink_util::store_trait::{ use nativelink_util::task::TaskExecutor; use nativelink_util::telemetry::init_tracing; use nativelink_util::{background_spawn, fs, spawn}; -use nativelink_worker::local_worker::new_local_worker; +use nativelink_worker::local_worker::{WorkerRegistration, new_local_worker}; use rustls_pki_types::pem::PemObject; use rustls_pki_types::{CertificateRevocationListDer, PrivateKeyDer}; use tokio::net::{TcpListener, TcpSocket}; @@ -87,7 +87,6 @@ static GLOBAL: MiMalloc = MiMalloc; const DEFAULT_ADMIN_API_PATH: &str = "/admin"; // Note: This must be kept in sync with the documentation in `HealthConfig::path`. -const DEFAULT_HEALTH_STATUS_CHECK_PATH: &str = "/status"; // Note: This must be kept in sync with the documentation in // `OriginEventsConfig::max_event_queue_size`. @@ -191,6 +190,16 @@ macro_rules! service_setup { }}; } +/// A JSON body with its content type, for the admin API's listings. +const fn json_response( + body: String, +) -> ([(axum::http::header::HeaderName, &'static str); 1], String) { + ( + [(axum::http::header::CONTENT_TYPE, "application/json")], + body, + ) +} + async fn inner_main( cfg: CasConfig, shutdown_tx: broadcast::Sender, @@ -205,6 +214,7 @@ async fn inner_main( let health_registry_builder = Arc::new(AsyncMutex::new(HealthRegistryBuilder::new("nativelink"))); + let mut worker_registrations: Vec> = Vec::new(); let store_manager = Arc::new(StoreManager::new()); { @@ -222,6 +232,24 @@ async fn inner_main( .err_tip(|| format!("Failed to add store '{name}'"))?; } store_manager.run_post_init().await?; + + // Workers start after the listeners are up, but the health registry + // is built with the listeners, so their registration flags are + // made and registered here and handed to the workers later. + let mut workers_health = health_registry_lock.sub_builder("workers"); + for (i, worker_cfg) in cfg.workers.iter().flatten().enumerate() { + let WorkerConfig::Local(local_worker_cfg) = worker_cfg; + let name = if local_worker_cfg.name.is_empty() { + format!("worker_{i}") + } else { + local_worker_cfg.name.clone() + }; + let registration = WorkerRegistration::new(&name); + workers_health + .sub_builder(&name) + .register_indicator(registration.clone()); + worker_registrations.push(registration); + } } let mut root_futures: Vec>> = Vec::new(); @@ -409,12 +437,15 @@ async fn inner_main( )); if let Some(health_cfg) = services.health { - let path = if health_cfg.path.is_empty() { - DEFAULT_HEALTH_STATUS_CHECK_PATH - } else { - &health_cfg.path - }; - svc = svc.route_service(path, HealthServer::new(health_registry, &health_cfg)); + let (path, readiness_path) = health_paths(&health_cfg)?; + svc = svc.route_service( + &path, + HealthServer::new(health_registry.clone(), &health_cfg), + ); + svc = svc.route_service( + &readiness_path, + HealthServer::readiness(health_registry, &health_cfg), + ); } if let Some(admin_config) = services.admin { @@ -424,9 +455,51 @@ async fn inner_main( &admin_config.path }; let worker_schedulers = Arc::new(worker_schedulers.clone()); + let listing_worker_schedulers = worker_schedulers.clone(); + let listing_action_schedulers = Arc::new(action_schedulers.clone()); + // Read-only views for the provisioner and operators: what the + // scheduler has connected and what it has queued, with the + // reservations the queued actions carry. + let admin_router = Router::new() + .route( + "/scheduler/{instance_name}/workers", + axum::routing::get(move |params: axum::extract::Path| async move { + let instance_name = params.0; + let Some(scheduler) = listing_worker_schedulers.get(&instance_name) else { + return Err(( + StatusCode::NOT_FOUND, + format!("No scheduler named '{instance_name}'"), + )); + }; + let workers = scheduler.worker_snapshot().await; + nativelink_scheduler::admin::workers_json(&workers) + .map(json_response) + .map_err(|e| { + (StatusCode::INTERNAL_SERVER_ERROR, format!("Error: {e:?}")) + }) + }), + ) + .route( + "/scheduler/{instance_name}/demand", + axum::routing::get(move |params: axum::extract::Path| async move { + let instance_name = params.0; + let Some(scheduler) = listing_action_schedulers.get(&instance_name) else { + return Err(( + StatusCode::NOT_FOUND, + format!("No scheduler named '{instance_name}'"), + )); + }; + nativelink_scheduler::admin::queued_demand_json(scheduler.as_ref()) + .await + .map(json_response) + .map_err(|e| { + (StatusCode::INTERNAL_SERVER_ERROR, format!("Error: {e:?}")) + }) + }), + ); svc = svc.nest_service( path, - Router::new().route( + admin_router.route( "/scheduler/{instance_name}/set_drain_worker/{worker_id}/{is_draining}", axum::routing::post( move |params: axum::extract::Path<(String, String, String)>| async move { @@ -742,7 +815,8 @@ async fn inner_main( historical_store, ) .await - .err_tip(|| "Could not make LocalWorker")?; + .err_tip(|| "Could not make LocalWorker")? + .with_registration(worker_registrations[i].clone()); let name = if local_worker.name().is_empty() { format!("worker_{i}") diff --git a/web/apps/docs/content/docs/configuration/scheduler-and-workers.mdx b/web/apps/docs/content/docs/configuration/scheduler-and-workers.mdx index dc3386d548a..5d96b620628 100644 --- a/web/apps/docs/content/docs/configuration/scheduler-and-workers.mdx +++ b/web/apps/docs/content/docs/configuration/scheduler-and-workers.mdx @@ -164,7 +164,13 @@ client gives up. **`max_job_retries`** (default 3): how many times an action that fails with an internal error is retried before the last error is returned to the client. It exists to stop one poisonous action from cycling through and destabilising -your whole fleet. +your whole fleet. A lost worker and a memory escalation do not spend it. + +**`max_worker_loss_retries`** (default 10): how many times an action whose +worker was lost (disconnected, timed out, evicted, or OOM-killed as a whole) +is queued again before the client gets a `FAILED_PRECONDITION`. Losing a +worker is not the action's failure, so this budget is separate and larger; +it only stops an action that takes a worker down every time it runs. **`retain_completed_for_s`** (default 60 s): how long a completed action's result stays available for a late `WaitExecution` call. Clients that @@ -175,6 +181,119 @@ On the worker side, **`max_inflight_tasks`** (default 0, meaning unlimited) is the one to set on any real machine. Unlimited means the worker accepts every action the scheduler offers and lets the OS sort out the contention. +**`drain_on_shutdown`** (off by default) makes a worker announce that it is +draining the moment it gets `SIGTERM`, so the scheduler stops dispatching to +it while it finishes what it holds. It needs a scheduler that understands the +drain flag; an older one removes the worker on the spot and requeues its +actions. + +Three bounds keep one action from taking the worker with it. +**`max_captured_output_bytes`** (default 0, unbounded) caps how much of an +action's stdout or stderr the worker holds in memory; past it the output +spills to a file under the action directory and is uploaded from disk, so a +tool that prints gigabytes costs disk, not RAM. **`precondition_timeout_ms`** +(default 30000) ends a hanging precondition script and refuses the action as +backpressure. **`orphan_sweep_interval_s`** (default 0, off) removes action +directories under `work_directory` that no running action owns and that a +failed cleanup left behind; the directory is always purged at startup. + +When the worker ends an action, for a timeout or a scheduler cancel, the +process group gets SIGTERM first and **`kill_grace_ms`** (default 5000) to +write its own cleanup before SIGKILL follows; the memory reservation kill +skips the grace. A killed action keeps the stdout and stderr it produced. +**`set_tmpdir`** (default true) gives every action its own `TMPDIR` under its +action directory, so concurrent actions writing the same temporary file name +no longer collide; a `TMPDIR` in the action's own environment still wins, and +code that hard-codes `/tmp` is unaffected. A timeout above `max_action_timeout_s` +is clamped to it rather than rejected. + +When the action's own process has exited, anything it left running in its +process group (a background job, a helper that daemonized) is killed, since +it would otherwise hold the action's output pipes open and the action would +sit there until its timeout; the pipes then get five seconds to drain, and +what was read by then is kept. The worker also makes itself the parent of +every process an action leaves behind and reaps them on a timer, so a build +whose actions fork and do not wait no longer leaves hundreds of zombies on +the worker for the life of the pod. Under `use_namespaces` the stub does +the reaping inside the action's PID namespace. + +Every executed action reports its resource usage to the scheduler with an +`outcome`: `COMPLETED`, `KILLED_MEMORY` (the worker's reservation kill, or a +SIGKILL from the kernel with the last memory sample within 10% of the limit), +`KILLED_TIMEOUT` or `KILLED_EXTERNAL`, plus wall time and the reservation the +action ran under. The scheduler logs every kill outcome with its mnemonic. + +**`resource_enforcement`** (off by default, Linux only) makes the worker hold an +action to the memory reservation the scheduler placed it by. With +`memory: soft` the worker kills the action's process group once two +consecutive samples exceed `memory_kb` plus `memory_headroom_percent` +(default 20), or one sample exceeds twice that, and fails it with +`FAILED_PRECONDITION` naming the reservation and the observed peak. Samples +are 250 ms apart, so an action that allocates faster than the pod's whole +limit in half a second can still reach the cgroup limit first; the +reservations admitted onto a worker must also sum to less than its limit +with the headroom included. Without it the kill comes from the pod's cgroup limit, +and on cgroup v2 that kill takes the worker and every action on it, not just +the one that grew. The sample is the group's proportional set size, so a +parent's pages shared with its forked children count once. Actions run +through a persistent worker are not covered: that process outlives the action +and serves others, so only the pod's limit bounds it. With enforcement on, the +worker reads an action's properties from the scheduler's dispatch rather than +the client's request, so a reservation the scheduler placed the action by is +the one enforced and exported to the environment. + +`disk: guard` adds a check before any input is fetched: an action whose +`disk_kb` reservation exceeds the free space under the work directory, less +what the actions already admitted reserved, is refused with `RESOURCE_EXHAUSTED`, which the scheduler requeues without +counting an attempt and holds off this worker until its next keepalive, +instead of running the fetch into a full disk. +`disk: soft` is the guard plus a kill: every ten seconds the worker adds up +the files the action has written under its own directory (those modified +since its command started; inputs were materialized before that and are +not counted), and an action past its `disk_kb` reservation by +`disk_headroom_percent` is killed and fails with `FAILED_PRECONDITION`, its +outcome `KILLED_DISK`. The walk runs on the blocking pool as its own task, +so memory sampling carries on meanwhile. With either disk mode the action's +usage report carries `peak_disk_kb`, what it wrote, from one walk of the +action directory at completion, which is the only walk `guard` pays. The +figure is reported and exported for now; the scheduler's sizing consumer +comes with the size ladder. + +```json5 +"resource_enforcement": { + "memory": "soft", + "memory_property_name": "memory_kb", + "memory_headroom_percent": 20, + "disk": "soft", + "disk_property_name": "disk_kb", + "disk_headroom_percent": 20 +} +``` + +**`capacity`** makes the worker advertise CPU and memory from its own cgroup +limits instead of from `platform_properties`: it reads `cpu.max` and +`memory.max` at its cgroup v2 root, keeps back `overhead_cpu_millicores` +(thousandths of a core) and `overhead_memory_kb` for itself, divides the +memory by the enforcement headroom (`memory_headroom_percent`, which follows +`resource_enforcement` unless set), and sets the two properties to the +result at registration. The limit is the number the kernel enforces, so the +number the scheduler packs against is derived from it rather than typed next +to it. Where the cgroup cannot be read, hand-typed `cpu_count` and +`memory_kb` stand with a warning; with neither typed the worker refuses to +start, since it could take no action that asks for either. `cpu_unit` is the scale of the CPU +property: `cores` by default, rounded down, the scale `cpu_count: { +query_cmd: "nproc" }` and an action's `cpu_count=1` use; a fleet that +advertises and requests thousandths of a core sets `millicores`. The +overhead is always in thousandths of a core. + +```json5 +"capacity": { + "cpu_unit": "cores", + "overhead_cpu_millicores": 1000, + "overhead_memory_kb": 4194304 +} +``` + `experimental_backend` defaults to memory, so restarting the scheduler loses the queue. A Redis backend exists for shared state across scheduler diff --git a/web/apps/docs/content/docs/contribute/docs.mdx b/web/apps/docs/content/docs/contribute/docs.mdx index 8a38b35b585..f9cf38aed0e 100644 --- a/web/apps/docs/content/docs/contribute/docs.mdx +++ b/web/apps/docs/content/docs/contribute/docs.mdx @@ -211,8 +211,10 @@ a how-to showing a renamed field keeps showing it, and nothing fails until a reader copies it and the binary answers `unknown field`. `lint:snippets` closes that gap. It harvests every field name out of the -generated reference and checks every `json`/`json5` fence under -`content/docs/` against it. +generated reference, the latest release's page and main's, and checks every +`json`/`json5` fence under `content/docs/` against it. A snippet for a field +added since the last release passes once `gen:config-reference main` has +been run in the same change. ```bash bun --filter @nativelink/docs lint:snippets diff --git a/web/apps/docs/content/docs/explanations/correctness-hermeticity.mdx b/web/apps/docs/content/docs/explanations/correctness-hermeticity.mdx index bf479a6e480..6357f171f57 100644 --- a/web/apps/docs/content/docs/explanations/correctness-hermeticity.mdx +++ b/web/apps/docs/content/docs/explanations/correctness-hermeticity.mdx @@ -99,8 +99,11 @@ is the toolchain half of the answer. ### Actions can run more than once [#actions-can-run-more-than-once] The default `max_job_retries` of 3 is checked strictly, so an action may -execute **four** times. A worker disconnect always costs an attempt; a -`ResourceExhausted` never does. +execute **four** times on its own failures. A lost worker has its own budget, +`max_worker_loss_retries` (default 10), and a memory escalation its own +(`memory_escalation.max_steps`, default 8), so an action can run more often +than four times when workers die under it or it keeps outgrowing its +reservation; a `ResourceExhausted` never costs anything. For a pure function this is invisible. For an action with side effects (one that publishes, deploys, increments, or writes anywhere outside its declared diff --git a/web/apps/docs/content/docs/explanations/scheduler-internals.mdx b/web/apps/docs/content/docs/explanations/scheduler-internals.mdx index e7db91bfa99..e8cf2c4c519 100644 --- a/web/apps/docs/content/docs/explanations/scheduler-internals.mdx +++ b/web/apps/docs/content/docs/explanations/scheduler-internals.mdx @@ -140,11 +140,16 @@ values per worker. For clusters of tens or low hundreds of workers this is not the bottleneck. It is worth knowing about before you plan for thousands. -`allocation_strategy` is a smaller knob than its name suggests: it chooses -**which end of the LRU the scan starts from**, and nothing else. Starting -from the least-recently-used end spreads load; starting from the -most-recently-used end packs work onto hot workers and leaves the rest idle -long enough to be scaled down. Neither is a bin-packing algorithm. +`allocation_strategy` is a smaller knob than its name suggests. The two +original values choose **which end of the LRU the scan starts from**, and +nothing else. Starting from the least-recently-used end spreads load; +starting from the most-recently-used end packs work onto hot workers and +leaves the rest idle long enough to be scaled down. `least_loaded` takes the +eligible worker with the fewest running actions, least recently used among +equals, so a worker that just joined takes the next action instead of +waiting for every older worker to be used once more, and a burst that lands +while most of the fleet is still starting does not fill the first worker to +its concurrency cap. None of them is a bin-packing algorithm. The admission ledger only subtracts what an action declares, so an action that declared too little is invisible to it. **`live_memory_veto`** names a @@ -162,6 +167,52 @@ a worker, which is what `worker_timeout_s` should be judged against, and Within a priority the queue is ordered by insert time to the nanosecond, on both backends, so a burst submitted together dispatches in arrival order. +The pass keeps that order when a worker frees up while it is running. A pass +reads the queue one action at a time, a store round trip each on a +store-backed queue (the listing serves the record it already loaded, and +only the client keepalive is read), so on a long queue a worker that +finishes mid-pass would otherwise hand its room to whatever action the pass +happened to be reading, newer than everything it had already passed over. Instead, actions that found +no room are parked in the order the queue listed them, highest priority first +and then oldest, and when the fleet's capacity changes the pass offers the +room to the parked actions in that order before the one in hand, asking the +fleet once per property shape. `scheduler.matching.parked_dispatched` counts +the parked actions a worker was actually told to run that way; a non-zero +rate is normal on a busy fleet and means the head of the queue is being +served. + +When a worker reports an action killed for memory, by its own reservation +enforcement or by the kernel near the limit, **`memory_escalation`** requeues +the action with the reservation multiplied by `percent` (default 200) instead +of failing it. With `ladder_kb` set to the memory values of the fleet's size +classes, a kill steps the reservation to the next class up instead of +scaling it. The next dispatch carries the raised value, so the worker holds +the action to that. The reservation grows to the largest memory any connected +worker advertises (or `max_kb`); the last step reserves that worker whole, +its memory and, through `cpu_property`, its CPU, so the action runs alone. +Only a kill there fails the action, with a message saying no worker in the +fleet can run it. Escalations have their own budget (`max_steps`, default 8) +and do not spend `max_job_retries`. The execution histograms carry an +`execution.outcome` attribute so the killed samples can be told from the +completed ones. + +The `historical_resource` scheduler that fills reservations reads hints by a +key ladder, each step probed only when the action carries that key: +`(target_id, action_mnemonic)`, then `target_id`, then the +action's `action_digest`, then its `command_digest` (each `hash-size` or +`hash/size`), then the mnemonic alone. The digest keys reach every client, +whether or not it sends Bazel request metadata; the action digest is what +origin events carry, so it is what a producer fed by them writes. A hint may name a `class` instead of numbers, and +**`classes`** on the same scheduler defines the fleet's size ladder: named +points with `cpu_count`, `memory_kb` and `disk_kb`, listed ascending. An +action no hint matches takes the class picked by `class_by_mnemonic`, then +`class_by_timeout` (Bazel test sizes arrive as timeouts), then +`default_class`; without classes the raw `cold_start` numbers apply. With +`memory_escalation.ladder_kb` set to the classes' memory values, a memory +kill steps the reservation to the next class instead of scaling it, and a +kill at the top class fails the action. Older hint files still load: every +new field is optional. + ### Platform properties are typed, and `Minimum` is a resource pool [#platform-properties-are-typed-and-minimum-is-a-resource-pool] Every property key a worker advertises has a declared type: @@ -197,17 +248,25 @@ the configuration side. ## Retries, and what "failed" means [#retries-and-what-failed-means] -The attempt counter is incremented by two different events, with different -rules. +The attempt counter is incremented by one event, and two other events that +used to spend it now have budgets of their own. **`UpdateWithError` increments the attempt count, unless the error is `ResourceExhausted`.** Backpressure is free. A worker saying "I am too busy" does not spend one of the action's lives, which is what makes it safe for a worker to shed load. -**`UpdateWithDisconnect` always increments.** A worker that vanishes mid-action -always costs an attempt. Without this, an action that reliably OOM-kills its -worker would retry forever, taking a worker down each time. +**`UpdateWithDisconnect` counts against `max_worker_loss_retries`** (default +10), not `max_job_retries`. A worker that vanishes mid-action is usually not +the action's doing, so the action goes back to the queue without spending one +of its lives; the separate budget is what stops an action that reliably +OOM-kills its worker from taking one down forever. Past it the client sees a +`FAILED_PRECONDITION` saying the worker was lost that many times. + +**`UpdateWithEscalation` counts against `memory_escalation.max_steps`** +(default 8), not `max_job_retries` either: the scheduler chose to run the +action again with more memory, and the fleet's ceiling bounds how far that +goes. The limit is checked as `attempts > max_job_retries`, strictly, so the default of `3` permits **four** total executions. That off-by-one is in the diff --git a/web/apps/docs/content/docs/explanations/worker-execution.mdx b/web/apps/docs/content/docs/explanations/worker-execution.mdx index 793a2694f51..d2d162697f7 100644 --- a/web/apps/docs/content/docs/explanations/worker-execution.mdx +++ b/web/apps/docs/content/docs/explanations/worker-execution.mdx @@ -178,7 +178,8 @@ command finished slightly before the bytes exist. participant sched as Scheduler participant w as Worker participant cas as CAS + AC - sched->>w: assign action (over ConnectWorker) + sched->>w: StartExecute (over ConnectWorker) + w->>sched: ExecuteAccepted, or ExecuteDeclined and stop here w->>w: create work directory for the operation w->>cas: fetch missing blobs w->>w: hardlink inputs from filesystem CAS @@ -190,6 +191,24 @@ command finished slightly before the bytes exist. w->>sched: final result update `} +The first reply is the worker's admission decision. A dispatch is a +proposal until the worker answers it: `ExecuteAccepted` confirms the ledger +charge the scheduler made on the send, `ExecuteDeclined` hands it back. The +worker declines when it already holds `max_inflight_tasks` actions, when the +action's memory reservation is more than the worker currently has free, or +when it is shutting down. A declined action returns to the queue without +counting as an attempt, and the worker is paused until its next keepalive, +or, for a decline for memory, until a keepalive reports that much free. The +scheduler's channel to each worker is bounded at its concurrency plus a +margin; a worker that has stopped reading has its dispatch requeued the same +way rather than piling up work it will never see. A worker only sends these +replies to a scheduler that announced support in its `ConnectionResult`, so +an older scheduler sees the old behaviour: the worker runs whatever it is +sent. `dispatch_ack_timeout_s` on the scheduler adds the converse: a dispatch +no worker acknowledged within the timeout is taken back, which only makes +sense once every connected worker acknowledges. + + Four behaviours in the upload path: **A declared-but-missing output is not an error.** An action that promises diff --git a/web/apps/docs/content/docs/how-to/stores/compose-stores.mdx b/web/apps/docs/content/docs/how-to/stores/compose-stores.mdx index e96b68d1f08..7681673d0ac 100644 --- a/web/apps/docs/content/docs/how-to/stores/compose-stores.mdx +++ b/web/apps/docs/content/docs/how-to/stores/compose-stores.mdx @@ -319,6 +319,23 @@ restart. forward selected headers from the incoming request. Use it to front an existing cache. +Development builds also support `grpc.store_type: "event_sink"` for forwarding +BEP and origin events to a dedicated receiver. This mode accepts string keys +with exact payload sizes and sends uncompressed `ByteStream.Write` requests. +The resource hash identifies the event key, while the size describes the event +bytes. Configure it only as an event publisher's store: reads, existence checks, +and digest-key writes are rejected. A normal CAS endpoint cannot receive these +writes. The receiver must acknowledge the full payload size only after its +durable append completes; NativeLink rejects a different committed size, and +its BEP response waits for that acknowledgement. A lost acknowledgement can +cause a retry after the receiver has already persisted an event. The receiver +must use the stable event key to detect and ignore duplicate events, including +retries that start a new upload; the upload UUID alone is not an event identity. +Origin publishing retains its existing bounded retry policy. This transport +provides neither exactly-once processing nor a durable publisher recovery ledger. + + + `noop` takes nothing at all and discards everything written to it, reporting every read as a miss. It is how you deliberately turn a half of something off. diff --git a/web/apps/docs/content/docs/operate/autoscaling-reference.mdx b/web/apps/docs/content/docs/operate/autoscaling-reference.mdx index 7c3361d7271..fc8a27c762c 100644 --- a/web/apps/docs/content/docs/operate/autoscaling-reference.mdx +++ b/web/apps/docs/content/docs/operate/autoscaling-reference.mdx @@ -166,6 +166,43 @@ each check) and `200` otherwise. "nothing objected," which for most compositions is vacuous. +## Readiness is a different question [#readiness-is-a-different-question] + +`/status` answers whether the components that registered a health indicator +have failed. It stays `200` while a component is still initializing, which is +the right answer for a liveness probe and the wrong one for readiness. The +`health` service therefore also serves `/ready` (`readiness_path`, default +`/ready`), which answers `503` while anything is initializing or failed and +`200` otherwise. + +For a worker, the component that matters is its registration with the +scheduler. Each worker registers a `workers/` indicator that reads +`Initializing` until the scheduler has answered its `connect_worker` with a +`ConnectionResult`, `Ok` while that connection holds, and `Initializing` +again after it is lost and the worker is reconnecting, so only `/ready` +drops: a liveness probe on `/status` stays green through a scheduler roll +instead of restarting every worker. A Kubernetes readiness probe on the +worker's `/ready` turns the pod Ready when the worker can take work rather +than when its container started, and turns it not Ready again if it loses +the scheduler. Anything that reads pod readiness, an autoscaler's idle clock +included, then counts from registration. + +```json5 +// On the worker pod's HTTP listener. +services: { + health: { path: "/status", readiness_path: "/ready" }, +} +``` + +```yaml +readinessProbe: + httpGet: { path: /ready, port: 50061 } + periodSeconds: 5 + failureThreshold: 3 +``` + + + ## Build the deployment [#build-the-deployment] diff --git a/web/apps/docs/content/docs/operate/deploy-bare-metal.mdx b/web/apps/docs/content/docs/operate/deploy-bare-metal.mdx index f5efdfb9a61..7340cfce55e 100644 --- a/web/apps/docs/content/docs/operate/deploy-bare-metal.mdx +++ b/web/apps/docs/content/docs/operate/deploy-bare-metal.mdx @@ -243,10 +243,14 @@ reconnect path calls `kill_all()` first, destroying every action it was running. Restart the scheduler when the queue is quiet, and restart it once. **Workers** drain themselves on `SIGTERM` (and on a first `SIGINT`; a second -`SIGINT` exits at once): the process rejects new actions, -waits for the ones it is running to finish, sends the scheduler a -`GoingAway`, and exits. `systemctl stop` is therefore safe as long as -`TimeoutStopSec` outlasts your longest action. To take a worker out of +`SIGINT` exits at once): with `drain_on_shutdown: true` +the process tells the scheduler at once that it is draining (a `GoingAway` +with the drain flag), so nothing new is dispatched to it, finishes the +actions it is running, and exits; the scheduler removes it when the stream +closes. Without the flag the worker finishes first and announces itself +last, and until then the scheduler keeps offering it work it refuses. The +wait is bounded by `max_action_timeout_s`. `systemctl stop` is therefore +safe as long as `TimeoutStopSec` outlasts your longest action. To take a worker out of rotation without stopping it, and if your scheduler configuration enables the `admin` service on the worker API listener: diff --git a/web/apps/docs/content/docs/operate/deploy-docker-compose.mdx b/web/apps/docs/content/docs/operate/deploy-docker-compose.mdx index 133bef3f418..4482325a989 100644 --- a/web/apps/docs/content/docs/operate/deploy-docker-compose.mdx +++ b/web/apps/docs/content/docs/operate/deploy-docker-compose.mdx @@ -168,8 +168,9 @@ nativelink_executor: which need `CLONE_NEWUSER`/`CLONE_NEWNS`, hence the flag. Note what this buys and what it does not: namespaces here are a hermeticity and process-reaping mechanism, not a security boundary. Actions still run as the worker's own uid -through an identity uid map, still see the host filesystem, and still have -network access to every NativeLink port. A privileged container running +through an identity uid map, still see the host filesystem apart from a +private `/tmp`, and still have network access to every NativeLink port. A +privileged container running arbitrary build actions is a lateral-movement primitive. [Security hardening](/operate/security-hardening#the-sandbox-is-not-a-security-boundary) covers what the sandbox actually does. diff --git a/web/apps/docs/content/docs/operate/observability.mdx b/web/apps/docs/content/docs/operate/observability.mdx index fc5971fba4d..75f65fb194b 100644 --- a/web/apps/docs/content/docs/operate/observability.mdx +++ b/web/apps/docs/content/docs/operate/observability.mdx @@ -295,7 +295,8 @@ each one decides where you have to point the collector: | Cache | `nativelink_cache_*` | Any process with a store wrapped in `cache_metrics` | Hit rate, operation latency, bytes in and out, per cache type | | Execution | `nativelink_execution_*` | The scheduler | Queue time, stage durations, active and completed counts, retries, and per-action CPU time and peak memory as reported by the worker that ran it | | Worker fleet | `nativelink_worker_*` | The scheduler | Workers connected, joins and departures by reason, keepalives, and how many are paused or draining | -| Scheduler | `nativelink_scheduler_matching_*` | The scheduler | How long each matching pass takes and whether it succeeded | +| Scheduler | `nativelink_scheduler_matching_*` | The scheduler | How long each matching pass takes, whether it succeeded, and how many waiting actions took room that opened mid-pass ahead of newer ones | +| Sizing | `nativelink_scheduler_hints_resolutions_total` | The scheduler, with `historical_resource` configured | Actions sized, by source: a hint (`target`, `action_digest`, `command_digest`, `mnemonic`), `cold_start`, or `none`; the `cold_start` share is what a hint feed is meant to drive down | | gRPC server | `nativelink_rpc_server_duration_*` | Every process, for every gRPC service it serves | Request latency by service, method and status code | | Store tier | `nativelink_store_tier_*` | Any process with a `fast_slow` store | Fast and slow tier hits, misses and stale reads, and bytes read from each tier | | Health | `nativelink_health_checks_*` | Every process, on each run of the health endpoint | Health check outcomes per component: `ok`, `initializing`, `warning`, `failed` or `timeout` | @@ -348,6 +349,16 @@ sum(rate(nativelink_execution_completed_count_total{execution_result="success"}[ sum(rate(nativelink_execution_completed_count_total[5m])) ``` +Actions the worker killed for memory, as a share of the sampled ones. The +per-action resource histograms carry `execution_outcome` (`completed`, +`killed_memory`, `killed_disk`), so a kill can be told from a completion and +the memory escalation rate read against it: + +```text +sum(rate(nativelink_execution_peak_memory_count{execution_outcome="killed_memory"}[5m])) / +sum(rate(nativelink_execution_peak_memory_count[5m])) +``` + Workers connected to each scheduler instance, and how many of them are paused or draining: @@ -475,10 +486,10 @@ not your pipeline. Queue time, not queue depth. `nativelink_execution_queue_time` is a real histogram emitted by the scheduler; a p95 above your tolerance means add - workers. No `nativelink_scheduler_queue_depth` series exists; the - closest thing is `nativelink_execution_active_count{execution_stage="queued"}`, - a gauge, and `nativelink_worker_connected_count` tells you how many workers - that queue is waiting on. Before you build an HPA on either, read the + workers. `nativelink_scheduler_queue_depth` is a gauge of what the last + matching pass read, `nativelink_execution_active_count{execution_stage="queued"}` + counts the same thing from the other side, and `nativelink_worker_connected_count` + tells you how many workers that queue is waiting on. Before you build an HPA on either, read the [autoscaling reference](/operate/autoscaling-reference#the-reaction-time-budget): the default one-minute export interval alone puts your reaction time in minutes, and CPU-based scaling is worse because workers are mostly I/O-bound. diff --git a/web/apps/docs/content/docs/operate/production-config.mdx b/web/apps/docs/content/docs/operate/production-config.mdx index 75600f21eeb..1905a460d3d 100644 --- a/web/apps/docs/content/docs/operate/production-config.mdx +++ b/web/apps/docs/content/docs/operate/production-config.mdx @@ -402,8 +402,11 @@ workers: [{ `use_namespaces` puts each action in its own process namespace; `use_mount_namespace` additionally isolates the worker root in a new mount -namespace, and only works when `use_namespaces` is on. Both are off by default -and both are recommended in production. The container +namespace and gives each action a private `/tmp` (its own temporary directory, +bound over `/tmp` and removed with the action), and only works when +`use_namespaces` is on. Both are off by default and both are recommended in +production. Set `isolate_tmp: false` if actions must see the host's `/tmp`, +for example because a tool they run lives there. The container running the worker needs the privileges to create namespaces; the reference Compose stack marks the executor service `privileged: true` for exactly this. diff --git a/web/apps/docs/content/docs/operate/runbooks.mdx b/web/apps/docs/content/docs/operate/runbooks.mdx index a529e70616b..3f2463be338 100644 --- a/web/apps/docs/content/docs/operate/runbooks.mdx +++ b/web/apps/docs/content/docs/operate/runbooks.mdx @@ -69,15 +69,16 @@ signal deaths, so the log line is what you alert on. **The worker process itself was killed.** The scheduler notices after `worker_timeout_s` (default 10 seconds, checked once a second), evicts it, and -re-queues its in-flight actions. In this tree a disconnect *counts as an -attempt*, so after `max_job_retries` the client finally sees, as a +re-queues its in-flight actions. A lost worker does not spend the action's +`max_job_retries`; it has its own budget, `max_worker_loss_retries` (default +10), so after that many losses the client finally sees, as a `FAILED_PRECONDITION` that Bazel does not resubmit: ```text -Worker disconnected repeatedly while executing this action (4 > 3 attempts); -the worker was likely OOM-killed or its pod evicted. Give the action a memory -reservation or raise the pool's memory limit before retrying. for operation_id: ..., -maybe_worker_id: ... +The worker running this action was lost 11 times, more than +max_worker_loss_retries (10); the action is likely what takes the worker down +(OOM-killed pod, eviction). Give it a memory reservation or raise the pool's +memory limit before retrying. for operation_id: ..., maybe_worker_id: ... ``` That string is the single best diagnostic for a memory-pressured worker fleet. @@ -410,6 +411,41 @@ operation whose record is gone, count it, and carry on, so an eviction no longer makes the pile permanent. It still makes every affected client wait out its timeout. +## The scheduler's Redis keeps every completed action [#scheduler-redis-keeps-completed-actions] + +**Symptom.** With the Redis scheduler backend, `DBSIZE` on the scheduler's +database grows by one `aa_*` hash per action ever run and never falls; the +records all read `completed` and `TTL` says `-1`. Once the search index holds +more than one cursor page, listings over it (the admin read the worker +provisioner makes, the scheduler admin endpoints) can fail with +`Failed to get operations list! Parse: Non map item`. + +**Cause.** Before v1.7.3 the Redis backend received `retain_completed_for_s` +as configured, without the default applied: left unset it arrived as 0, and the versioned update treats a +zero TTL as no expiry, so completed records were written to live forever. +The memory backend applied the 60 second default; only Redis leaked. The +listing failure is a second bug in the same release: a cursor page arrived +in the RESP3 shape the first page's parser understood but the cursor reader +did not. + +**Fix.** Upgrade to v1.7.3 or later; new completions expire after +`retain_completed_for_s`. Records written before the upgrade keep their +missing expiry, so give them one once, on the scheduler's Redis database +(0 in the chart's layout): + +```bash +redis-cli -n 0 --scan --pattern 'aa_*' \ + | while read -r key; do + if [ "$(redis-cli -n 0 hget "$key" state)" = "completed" ] \ + && [ "$(redis-cli -n 0 ttl "$key")" = "-1" ]; then + redis-cli -n 0 expire "$key" 60 > /dev/null + fi + done +``` + +A queued or executing record is left alone; only completed ones get the +sixty seconds a fresh completion would. + ## Restarting the scheduler [#restarting-the-scheduler] Not an incident so much as a procedure that behaves worse than you expect. diff --git a/web/apps/docs/content/docs/operate/scaling-workers.mdx b/web/apps/docs/content/docs/operate/scaling-workers.mdx index 82057cde28b..1605ddaefec 100644 --- a/web/apps/docs/content/docs/operate/scaling-workers.mdx +++ b/web/apps/docs/content/docs/operate/scaling-workers.mdx @@ -247,6 +247,22 @@ exist at all. Nothing in the repository wires this up for you. It is a short script, and it's the difference between a clean scale-in and a burst of retried actions. +Two read-only listings sit next to it, for the same scripts and for an +autoscaler that wants the scheduler's own view rather than its store: + +``` +GET {admin_path}/scheduler/{scheduler_name}/workers +GET {admin_path}/scheduler/{scheduler_name}/demand +``` + +`workers` returns every connected worker with what it registered, what it +has left after its running actions' reservations, its running count, its +pause and drain flags and the free memory it last reported. `demand` +returns every queued action with its operation id, when it was queued, its +priority and its platform properties, reservations included. A worker's id +starts with the name it registered under, which on Kubernetes is the pod +name, so a drain script can find the ids to drain from the pod. +
    • diff --git a/web/apps/docs/content/docs/operate/security-hardening.mdx b/web/apps/docs/content/docs/operate/security-hardening.mdx index 5bf06044f7b..e7c098ca64d 100644 --- a/web/apps/docs/content/docs/operate/security-hardening.mdx +++ b/web/apps/docs/content/docs/operate/security-hardening.mdx @@ -319,9 +319,14 @@ uid-remapping sandbox; it is a process-reaping and hermeticity mechanism. PID namespacing means orphaned children die with the action instead of leaking; IPC and UTS namespacing means actions cannot see each other's semaphores or hostname (the hostname is set to `nativelink`). With `use_mount_namespace` -the flag list also gains `CLONE_NEWNS`, and the worker's own root action -directory is masked so one action cannot see another's tree; the rest of the -host filesystem stays visible. +the flag list also gains `CLONE_NEWNS`, the worker's own root action +directory is masked so one action cannot see another's tree, and `/tmp` is +replaced with the action's own temporary directory, on the worker's disk and +removed with the action, so actions cannot collide on, or leak state +through, predictable paths there (`isolate_tmp: false` keeps the host's +`/tmp` visible; the bound `/tmp` carries the work volume's mount options, +`noexec` included, not the host `/tmp`'s); the rest of the host filesystem +stays visible. Note what is absent from that flag list: **`CLONE_NEWNET`**. No network namespace is created. An action has the worker's full network access, including to port diff --git a/web/apps/docs/content/docs/reference/metrics.mdx b/web/apps/docs/content/docs/reference/metrics.mdx index 28ec1584121..3b40409eb09 100644 --- a/web/apps/docs/content/docs/reference/metrics.mdx +++ b/web/apps/docs/content/docs/reference/metrics.mdx @@ -5,12 +5,12 @@ full: true --- {/* AUTOGENERATED. Do not edit by hand. - Source: nativelink-util/src/metrics.rs @ v1.6.5 (c479a38f) + Source: nativelink-util/src/metrics.rs @ v1.7.2 (838aeade6) Regenerate from web/: bun --filter @nativelink/docs gen:metrics-reference */} -NativeLink **v1.6.5** declares 30 OpenTelemetry instruments. This page is +NativeLink **v1.7.2** declares 40 OpenTelemetry instruments. This page is generated from the declarations in -[`nativelink-util/src/metrics.rs`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs) +[`nativelink-util/src/metrics.rs`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs) and from every `.add()` and `.record()` call site across the workspace, so the **Emitted** column is not a claim, it is the result of looking. @@ -20,15 +20,6 @@ series when some code path calls it. A catalogue written by hand cannot tell those two states apart, and the gap between them is where dashboards go quietly wrong. - - - They are constructed at startup and stored on the metrics struct, but no - code path calls them. See [Declared but never - emitted](#declared-but-never-emitted) for which ones, and for the shipped - Prometheus rules that query them anyway. - - - ## Every instrument at a glance | Instrument | Prometheus series | Kind | Emitted | Under what conditions | @@ -36,8 +27,8 @@ wrong. | [`cache.operation.duration`](#cache-operation-duration) | `nativelink_cache_operation_duration` | Histogram | Conditionally | only for stores explicitly wrapped in a `cache_metrics` spec | | [`cache.operations`](#cache-operations) | `nativelink_cache_operations_total` | Counter | Conditionally | only for stores explicitly wrapped in a `cache_metrics` spec | | [`cache.io`](#cache-io) | `nativelink_cache_io_total` | Counter | Conditionally | only for stores explicitly wrapped in a `cache_metrics` spec | -| [`cache.size`](#cache-size) | `nativelink_cache_size` | UpDownCounter | **Never** | no call site anywhere in the workspace | -| [`cache.entries`](#cache-entries) | `nativelink_cache_entries` | UpDownCounter | **Never** | no call site anywhere in the workspace | +| [`cache.size`](#cache-size) | `nativelink_cache_size` | UpDownCounter | Yes | whenever the code path runs | +| [`cache.entries`](#cache-entries) | `nativelink_cache_entries` | UpDownCounter | Yes | whenever the code path runs | | [`cache.item.size`](#cache-item-size) | `nativelink_cache_item_size` | Histogram | Conditionally | only for stores explicitly wrapped in a `cache_metrics` spec | | [`execution.stage.duration`](#execution-stage-duration) | `nativelink_execution_stage_duration` | Histogram | Conditionally | scheduler process only | | [`execution.total.duration`](#execution-total-duration) | `nativelink_execution_total_duration` | Histogram | Conditionally | scheduler process only | @@ -52,12 +43,21 @@ wrong. | [`worker.connected.count`](#worker-connected-count) | `nativelink_worker_connected_count` (unreferenced) | UpDownCounter | Conditionally | scheduler process only | | [`worker.connections`](#worker-connections) | `nativelink_worker_connections` (unreferenced) | Counter | Conditionally | scheduler process only | | [`worker.disconnections`](#worker-disconnections) | `nativelink_worker_disconnections` (unreferenced) | Counter | Conditionally | scheduler process only | +| [`scheduler.dispatch.requeues`](#scheduler-dispatch-requeues) | `nativelink_scheduler_dispatch_requeues` (unreferenced) | Counter | Conditionally | scheduler process only | | [`worker.keepalives`](#worker-keepalives) | `nativelink_worker_keepalives` (unreferenced) | Counter | Conditionally | scheduler process only | | [`worker.keepalive.gap`](#worker-keepalive-gap) | `nativelink_worker_keepalive_gap` (unreferenced) | Histogram | Conditionally | scheduler process only | | [`worker.state.count`](#worker-state-count) | `nativelink_worker_state_count` (unreferenced) | UpDownCounter | Conditionally | scheduler process only | | [`rpc.server.duration`](#rpc-server-duration) | `nativelink_rpc_server_duration` (unreferenced) | Histogram | Yes | whenever the code path runs | | [`scheduler.matching.duration`](#scheduler-matching-duration) | `nativelink_scheduler_matching_duration` (unreferenced) | Histogram | Conditionally | scheduler process only | | [`scheduler.matching.passes`](#scheduler-matching-passes) | `nativelink_scheduler_matching_passes` (unreferenced) | Counter | Conditionally | scheduler process only | +| [`scheduler.unsatisfiable.queued`](#scheduler-unsatisfiable-queued) | `nativelink_scheduler_unsatisfiable_queued` (unreferenced) | Gauge | Conditionally | scheduler process only | +| [`scheduler.unsatisfiable.failed`](#scheduler-unsatisfiable-failed) | `nativelink_scheduler_unsatisfiable_failed` (unreferenced) | Counter | Conditionally | scheduler process only | +| [`scheduler.queue.depth`](#scheduler-queue-depth) | `nativelink_scheduler_queue_depth` (unreferenced) | Gauge | Conditionally | scheduler process only | +| [`scheduler.queue.retired`](#scheduler-queue-retired) | `nativelink_scheduler_queue_retired` (unreferenced) | Counter | Conditionally | scheduler process only | +| [`scheduler.matching.parked_dispatched`](#scheduler-matching-parked-dispatched) | `nativelink_scheduler_matching_parked_dispatched` (unreferenced) | Counter | Conditionally | scheduler process only | +| [`scheduler.awaited_action.orphans`](#scheduler-awaited-action-orphans) | `nativelink_scheduler_awaited_action_orphans` (unreferenced) | Counter | Conditionally | scheduler process only | +| [`scheduler.sweep.failures`](#scheduler-sweep-failures) | `nativelink_scheduler_sweep_failures` (unreferenced) | Counter | Conditionally | scheduler process only | +| [`scheduler.hints.resolutions`](#scheduler-hints-resolutions) | `nativelink_scheduler_hints_resolutions` (unreferenced) | Counter | Conditionally | scheduler process only | | [`store.tier.operations`](#store-tier-operations) | `nativelink_store_tier_operations` (unreferenced) | Counter | Conditionally | only through a `fast_slow` store | | [`store.tier.io`](#store-tier-io) | `nativelink_store_tier_io` (unreferenced) | Counter | Conditionally | only through a `fast_slow` store | | [`health.checks`](#health-checks) | `nativelink_health_checks` (unreferenced) | Counter | Yes | whenever the code path runs | @@ -72,7 +72,7 @@ binary. The names in the middle column above are what the shipped OpenTelemetry collector produces after two transformations: 1. The Prometheus exporter in - [`otel-collector-config.yaml`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/deployment-examples/metrics/otel-collector-config.yaml) + [`otel-collector-config.yaml`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/deployment-examples/metrics/otel-collector-config.yaml) is configured with `namespace: nativelink`, which prefixes every series with `nativelink_`. 2. The exporter rewrites the OTLP name into Prometheus form: dots become @@ -93,32 +93,13 @@ nothing has confirmed the exporter's spelling of it. ## Declared but never emitted -2 instruments are built by the -meter and stored on the metrics struct, but no code anywhere in the workspace -ever calls `.add()` or `.record()` on them. They will never appear in your -Prometheus output, at any configuration, on any release this page was -generated from. - -| Instrument | Series it would produce | Referenced by shipped config | -| --- | --- | --- | -| `cache.size` | `nativelink_cache_size` | [`deployment-examples/metrics/README.md:112`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/deployment-examples/metrics/README.md#L112), [`deployment-examples/metrics/prometheus-config.yaml:135`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/deployment-examples/metrics/prometheus-config.yaml#L135), [`deployment-examples/metrics/prometheus-recording-rules.yml:164`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/deployment-examples/metrics/prometheus-recording-rules.yml#L164) | -| `cache.entries` | `nativelink_cache_entries` | [`deployment-examples/metrics/README.md:113`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/deployment-examples/metrics/README.md#L113), [`deployment-examples/metrics/prometheus-recording-rules.yml:168`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/deployment-examples/metrics/prometheus-recording-rules.yml#L168) | - - - - `nativelink_cache_size` and `nativelink_cache_entries` appear in the recording rules and the metrics - README under `deployment-examples/metrics/`, but nothing emits them. Any - recording rule, dashboard panel or alert built on them evaluates to an - empty vector forever, which reads on a dashboard as a healthy zero, not - as a missing signal. Delete those panels or treat them as known-empty. - - +Every instrument declared in `metrics.rs` has at least one call site. ## Cache metrics Global cache metrics instruments. -Declared as `CACHE_METRICS` in [`nativelink-util/src/metrics.rs:419`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L419), backed by the `CacheMetrics` struct. +Declared as `CACHE_METRICS` in [`nativelink-util/src/metrics.rs:438`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L438), backed by the `CacheMetrics` struct. Every cache instrument is emitted from one place: the `cache_metrics` store wrapper. A store that is not wrapped emits nothing, no matter how the collector is configured. @@ -133,8 +114,8 @@ Histogram of cache operation durations in milliseconds | Unit | `ms` | | Prometheus series | `nativelink_cache_operation_duration` | | Emitted | Conditionally. only for stores explicitly wrapped in a `cache_metrics` spec | -| Declared at | [`nativelink-util/src/metrics.rs:423`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L423) | -| Emitted from | [`nativelink-store/src/cache_metrics_store.rs:67`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/cache_metrics_store.rs#L67) | +| Declared at | [`nativelink-util/src/metrics.rs:442`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L442) | +| Emitted from | [`nativelink-store/src/cache_metrics_store.rs:73`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/cache_metrics_store.rs#L73) | **Histogram buckets:** 19 explicit boundaries, from `0.001` to `5000.0` `ms`. Anything past the last boundary lands in the @@ -152,8 +133,8 @@ Counter of cache operations by type and result | Unit | — | | Prometheus series | `nativelink_cache_operations_total` | | Emitted | Conditionally. only for stores explicitly wrapped in a `cache_metrics` spec | -| Declared at | [`nativelink-util/src/metrics.rs:457`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L457) | -| Emitted from | [`nativelink-store/src/cache_metrics_store.rs:121`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/cache_metrics_store.rs#L121), [`nativelink-store/src/cache_metrics_store.rs:126`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/cache_metrics_store.rs#L126), [`nativelink-store/src/cache_metrics_store.rs:137`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/cache_metrics_store.rs#L137), [`nativelink-store/src/cache_metrics_store.rs:162`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/cache_metrics_store.rs#L162), [`nativelink-store/src/cache_metrics_store.rs:168`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/cache_metrics_store.rs#L168), [`nativelink-store/src/cache_metrics_store.rs:194`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/cache_metrics_store.rs#L194), [`nativelink-store/src/cache_metrics_store.rs:200`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/cache_metrics_store.rs#L200), [`nativelink-store/src/cache_metrics_store.rs:213`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/cache_metrics_store.rs#L213), [`nativelink-store/src/cache_metrics_store.rs:219`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/cache_metrics_store.rs#L219), [`nativelink-store/src/cache_metrics_store.rs:240`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/cache_metrics_store.rs#L240), [`nativelink-store/src/cache_metrics_store.rs:248`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/cache_metrics_store.rs#L248), [`nativelink-store/src/cache_metrics_store.rs:254`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/cache_metrics_store.rs#L254) | +| Declared at | [`nativelink-util/src/metrics.rs:476`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L476) | +| Emitted from | [`nativelink-store/src/cache_metrics_store.rs:127`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/cache_metrics_store.rs#L127), [`nativelink-store/src/cache_metrics_store.rs:132`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/cache_metrics_store.rs#L132), [`nativelink-store/src/cache_metrics_store.rs:143`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/cache_metrics_store.rs#L143), [`nativelink-store/src/cache_metrics_store.rs:168`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/cache_metrics_store.rs#L168), [`nativelink-store/src/cache_metrics_store.rs:174`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/cache_metrics_store.rs#L174), [`nativelink-store/src/cache_metrics_store.rs:200`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/cache_metrics_store.rs#L200), [`nativelink-store/src/cache_metrics_store.rs:206`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/cache_metrics_store.rs#L206), [`nativelink-store/src/cache_metrics_store.rs:219`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/cache_metrics_store.rs#L219), [`nativelink-store/src/cache_metrics_store.rs:225`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/cache_metrics_store.rs#L225), [`nativelink-store/src/cache_metrics_store.rs:246`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/cache_metrics_store.rs#L246), [`nativelink-store/src/cache_metrics_store.rs:254`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/cache_metrics_store.rs#L254), [`nativelink-store/src/cache_metrics_store.rs:260`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/cache_metrics_store.rs#L260) | ### `cache.io` @@ -166,8 +147,8 @@ Counter of bytes read/written during cache operations | Unit | `By` | | Prometheus series | `nativelink_cache_io_total` | | Emitted | Conditionally. only for stores explicitly wrapped in a `cache_metrics` spec | -| Declared at | [`nativelink-util/src/metrics.rs:462`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L462) | -| Emitted from | [`nativelink-store/src/cache_metrics_store.rs:74`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/cache_metrics_store.rs#L74), [`nativelink-store/src/cache_metrics_store.rs:242`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/cache_metrics_store.rs#L242) | +| Declared at | [`nativelink-util/src/metrics.rs:481`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L481) | +| Emitted from | [`nativelink-store/src/cache_metrics_store.rs:80`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/cache_metrics_store.rs#L80), [`nativelink-store/src/cache_metrics_store.rs:248`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/cache_metrics_store.rs#L248) | ### `cache.size` @@ -179,9 +160,9 @@ Current total size of all cached data in bytes | Field | `cache_size` | | Unit | `By` | | Prometheus series | `nativelink_cache_size` | -| Emitted | **Never**. no call site anywhere in the workspace | -| Declared at | [`nativelink-util/src/metrics.rs:468`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L468) | -| Emitted from | — | +| Emitted | Yes. whenever the code path runs | +| Declared at | [`nativelink-util/src/metrics.rs:487`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L487) | +| Emitted from | [`nativelink-util/src/evicting_map.rs:468`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/evicting_map.rs#L468) | ### `cache.entries` @@ -193,9 +174,9 @@ Current number of entries in cache | Field | `cache_entries` | | Unit | `{entry}` | | Prometheus series | `nativelink_cache_entries` | -| Emitted | **Never**. no call site anywhere in the workspace | -| Declared at | [`nativelink-util/src/metrics.rs:474`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L474) | -| Emitted from | — | +| Emitted | Yes. whenever the code path runs | +| Declared at | [`nativelink-util/src/metrics.rs:493`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L493) | +| Emitted from | [`nativelink-util/src/evicting_map.rs:468`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/evicting_map.rs#L468) | ### `cache.item.size` @@ -208,14 +189,14 @@ Histogram of individual cache entry sizes in bytes | Unit | `By` | | Prometheus series | `nativelink_cache_item_size` | | Emitted | Conditionally. only for stores explicitly wrapped in a `cache_metrics` spec | -| Declared at | [`nativelink-util/src/metrics.rs:480`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L480) | -| Emitted from | [`nativelink-store/src/cache_metrics_store.rs:77`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/cache_metrics_store.rs#L77) | +| Declared at | [`nativelink-util/src/metrics.rs:499`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L499) | +| Emitted from | [`nativelink-store/src/cache_metrics_store.rs:83`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/cache_metrics_store.rs#L83) | ## Execution metrics Global remote execution metrics instruments. -Declared as `EXECUTION_METRICS` in [`nativelink-util/src/metrics.rs:506`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L506), backed by the `ExecutionMetrics` struct. +Declared as `EXECUTION_METRICS` in [`nativelink-util/src/metrics.rs:539`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L539), backed by the `ExecutionMetrics` struct. Execution instruments are emitted from the scheduler process as actions change state, plus the per-action CPU time and peak memory a worker reports back when it finishes. Worker processes emit no execution metrics of their own. @@ -230,8 +211,8 @@ Histogram of stage durations in seconds | Unit | `s` | | Prometheus series | `nativelink_execution_stage_duration` | | Emitted | Conditionally. scheduler process only | -| Declared at | [`nativelink-util/src/metrics.rs:510`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L510) | -| Emitted from | [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:919`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L919) | +| Declared at | [`nativelink-util/src/metrics.rs:543`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L543) | +| Emitted from | [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:109`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L109) | **Histogram buckets:** 15 explicit boundaries, from `0.001` to `3600.0` `s`. Anything past the last boundary lands in the @@ -249,8 +230,8 @@ Histogram of total execution durations in seconds | Unit | `s` | | Prometheus series | `nativelink_execution_total_duration` | | Emitted | Conditionally. scheduler process only | -| Declared at | [`nativelink-util/src/metrics.rs:535`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L535) | -| Emitted from | [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:919`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L919) | +| Declared at | [`nativelink-util/src/metrics.rs:568`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L568) | +| Emitted from | [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:109`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L109) | **Histogram buckets:** 13 explicit boundaries, from `0.01` to `7200.0` `s`. Anything past the last boundary lands in the @@ -268,8 +249,8 @@ Histogram of queue wait times in seconds | Unit | `s` | | Prometheus series | `nativelink_execution_queue_time` | | Emitted | Conditionally. scheduler process only | -| Declared at | [`nativelink-util/src/metrics.rs:560`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L560) | -| Emitted from | [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:919`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L919) | +| Declared at | [`nativelink-util/src/metrics.rs:593`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L593) | +| Emitted from | [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:109`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L109) | **Histogram buckets:** 12 explicit boundaries, from `0.001` to `600.0` `s`. Anything past the last boundary lands in the @@ -287,8 +268,8 @@ Current number of actions in each stage | Unit | `{action}` | | Prometheus series | `nativelink_execution_active_count` (not referenced by any shipped rule or dashboard) | | Emitted | Conditionally. scheduler process only | -| Declared at | [`nativelink-util/src/metrics.rs:580`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L580) | -| Emitted from | [`nativelink-scheduler/src/memory_awaited_action_db.rs:662`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/memory_awaited_action_db.rs#L662), [`nativelink-scheduler/src/memory_awaited_action_db.rs:669`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/memory_awaited_action_db.rs#L669), [`nativelink-scheduler/src/memory_awaited_action_db.rs:814`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/memory_awaited_action_db.rs#L814) | +| Declared at | [`nativelink-util/src/metrics.rs:613`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L613) | +| Emitted from | [`nativelink-scheduler/src/memory_awaited_action_db.rs:425`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/memory_awaited_action_db.rs#L425), [`nativelink-scheduler/src/memory_awaited_action_db.rs:672`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/memory_awaited_action_db.rs#L672), [`nativelink-scheduler/src/memory_awaited_action_db.rs:679`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/memory_awaited_action_db.rs#L679), [`nativelink-scheduler/src/memory_awaited_action_db.rs:824`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/memory_awaited_action_db.rs#L824), [`nativelink-scheduler/src/store_awaited_action_db.rs:716`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/store_awaited_action_db.rs#L716) | ### `execution.completed.count` @@ -301,8 +282,8 @@ Total number of completed executions | Unit | `{action}` | | Prometheus series | `nativelink_execution_completed_count_total` | | Emitted | Conditionally. scheduler process only | -| Declared at | [`nativelink-util/src/metrics.rs:586`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L586) | -| Emitted from | [`nativelink-scheduler/src/memory_awaited_action_db.rs:688`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/memory_awaited_action_db.rs#L688), [`nativelink-scheduler/src/memory_awaited_action_db.rs:700`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/memory_awaited_action_db.rs#L700), [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:918`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L918), [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:928`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L928) | +| Declared at | [`nativelink-util/src/metrics.rs:619`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L619) | +| Emitted from | [`nativelink-scheduler/src/memory_awaited_action_db.rs:698`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/memory_awaited_action_db.rs#L698), [`nativelink-scheduler/src/memory_awaited_action_db.rs:710`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/memory_awaited_action_db.rs#L710), [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:108`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L108), [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:118`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L118) | ### `execution.stage.transitions` @@ -315,8 +296,8 @@ Number of stage transitions | Unit | `{transition}` | | Prometheus series | `nativelink_execution_stage_transitions_total` | | Emitted | Conditionally. scheduler process only | -| Declared at | [`nativelink-util/src/metrics.rs:592`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L592) | -| Emitted from | [`nativelink-scheduler/src/memory_awaited_action_db.rs:655`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/memory_awaited_action_db.rs#L655), [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:907`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L907) | +| Declared at | [`nativelink-util/src/metrics.rs:625`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L625) | +| Emitted from | [`nativelink-scheduler/src/memory_awaited_action_db.rs:665`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/memory_awaited_action_db.rs#L665), [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:97`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L97) | ### `execution.cpu.time` @@ -329,8 +310,8 @@ CPU time an action used, as sampled by the worker. | Unit | `s` | | Prometheus series | `nativelink_execution_cpu_time` (not referenced by any shipped rule or dashboard) | | Emitted | Conditionally. scheduler process only | -| Declared at | [`nativelink-util/src/metrics.rs:598`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L598) | -| Emitted from | [`nativelink-scheduler/src/api_worker_scheduler.rs:709`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/api_worker_scheduler.rs#L709) | +| Declared at | [`nativelink-util/src/metrics.rs:631`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L631) | +| Emitted from | [`nativelink-scheduler/src/api_worker_scheduler.rs:1784`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/api_worker_scheduler.rs#L1784) | **Histogram buckets:** 12 explicit boundaries, from `0.1` to `7200.0` `s`. Anything past the last boundary lands in the @@ -348,8 +329,8 @@ Peak memory an action used, as sampled by the worker. | Unit | `By` | | Prometheus series | `nativelink_execution_peak_memory` (not referenced by any shipped rule or dashboard) | | Emitted | Conditionally. scheduler process only | -| Declared at | [`nativelink-util/src/metrics.rs:607`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L607) | -| Emitted from | [`nativelink-scheduler/src/api_worker_scheduler.rs:706`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/api_worker_scheduler.rs#L706) | +| Declared at | [`nativelink-util/src/metrics.rs:640`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L640) | +| Emitted from | [`nativelink-scheduler/src/api_worker_scheduler.rs:1776`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/api_worker_scheduler.rs#L1776) | **Histogram buckets:** 8 explicit boundaries, from `1_048_576.0` to `68_719_476_736.0` `By`. Anything past the last boundary lands in the @@ -367,8 +348,8 @@ Histogram of output sizes in bytes | Unit | `By` | | Prometheus series | `nativelink_execution_output_size` | | Emitted | Conditionally. scheduler process only | -| Declared at | [`nativelink-util/src/metrics.rs:623`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L623) | -| Emitted from | [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:919`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L919) | +| Declared at | [`nativelink-util/src/metrics.rs:656`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L656) | +| Emitted from | [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:109`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L109) | **Histogram buckets:** 8 explicit boundaries, from `1_024.0` to `10_737_418_240.0` `By`. Anything past the last boundary lands in the @@ -386,14 +367,14 @@ Counter for execution retries | Unit | `{retry}` | | Prometheus series | `nativelink_execution_retry_count_total` | | Emitted | Conditionally. scheduler process only | -| Declared at | [`nativelink-util/src/metrics.rs:639`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L639) | -| Emitted from | [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:940`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L940) | +| Declared at | [`nativelink-util/src/metrics.rs:672`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L672) | +| Emitted from | [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:130`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L130) | ## Worker fleet metrics Global worker fleet metrics instruments. -Declared as `WORKER_METRICS` in [`nativelink-util/src/metrics.rs:817`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L817), backed by the `WorkerMetrics` struct. +Declared as `WORKER_METRICS` in [`nativelink-util/src/metrics.rs:881`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L881), backed by the `WorkerMetrics` struct. Worker fleet instruments are emitted from the scheduler process as workers connect, heartbeat, pause, drain and leave. They describe the pool as this scheduler instance sees it. @@ -408,8 +389,8 @@ Workers currently connected. | Unit | `{worker}` | | Prometheus series | `nativelink_worker_connected_count` (not referenced by any shipped rule or dashboard) | | Emitted | Conditionally. scheduler process only | -| Declared at | [`nativelink-util/src/metrics.rs:821`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L821) | -| Emitted from | [`nativelink-scheduler/src/api_worker_scheduler.rs:227`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/api_worker_scheduler.rs#L227), [`nativelink-scheduler/src/api_worker_scheduler.rs:504`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/api_worker_scheduler.rs#L504) | +| Declared at | [`nativelink-util/src/metrics.rs:885`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L885) | +| Emitted from | [`nativelink-scheduler/src/api_worker_scheduler.rs:442`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/api_worker_scheduler.rs#L442), [`nativelink-scheduler/src/api_worker_scheduler.rs:1378`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/api_worker_scheduler.rs#L1378) | ### `worker.connections` @@ -422,8 +403,8 @@ Workers that have joined, cumulative. | Unit | `{worker}` | | Prometheus series | `nativelink_worker_connections` (not referenced by any shipped rule or dashboard) | | Emitted | Conditionally. scheduler process only | -| Declared at | [`nativelink-util/src/metrics.rs:827`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L827) | -| Emitted from | [`nativelink-scheduler/src/api_worker_scheduler.rs:227`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/api_worker_scheduler.rs#L227) | +| Declared at | [`nativelink-util/src/metrics.rs:891`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L891) | +| Emitted from | [`nativelink-scheduler/src/api_worker_scheduler.rs:442`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/api_worker_scheduler.rs#L442) | ### `worker.disconnections` @@ -436,8 +417,22 @@ Workers that have left, cumulative, by reason. | Unit | `{worker}` | | Prometheus series | `nativelink_worker_disconnections` (not referenced by any shipped rule or dashboard) | | Emitted | Conditionally. scheduler process only | -| Declared at | [`nativelink-util/src/metrics.rs:833`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L833) | -| Emitted from | [`nativelink-scheduler/src/api_worker_scheduler.rs:504`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/api_worker_scheduler.rs#L504) | +| Declared at | [`nativelink-util/src/metrics.rs:897`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L897) | +| Emitted from | [`nativelink-scheduler/src/api_worker_scheduler.rs:1378`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/api_worker_scheduler.rs#L1378) | + +### `scheduler.dispatch.requeues` + +Dispatches requeued untried, by reason. + +| Property | Value | +| --- | --- | +| Instrument type | `Counter` | +| Field | `dispatch_requeues` | +| Unit | — | +| Prometheus series | `nativelink_scheduler_dispatch_requeues` (not referenced by any shipped rule or dashboard) | +| Emitted | Conditionally. scheduler process only | +| Declared at | [`nativelink-util/src/metrics.rs:903`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L903) | +| Emitted from | [`nativelink-scheduler/src/api_worker_scheduler.rs:1137`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/api_worker_scheduler.rs#L1137), [`nativelink-scheduler/src/api_worker_scheduler.rs:1574`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/api_worker_scheduler.rs#L1574), [`nativelink-scheduler/src/api_worker_scheduler.rs:1872`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/api_worker_scheduler.rs#L1872), [`nativelink-scheduler/src/api_worker_scheduler.rs:1903`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/api_worker_scheduler.rs#L1903) | ### `worker.keepalives` @@ -450,25 +445,27 @@ Keepalives received, cumulative. | Unit | `{keepalive}` | | Prometheus series | `nativelink_worker_keepalives` (not referenced by any shipped rule or dashboard) | | Emitted | Conditionally. scheduler process only | -| Declared at | [`nativelink-util/src/metrics.rs:839`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L839) | -| Emitted from | [`nativelink-scheduler/src/api_worker_scheduler.rs:187`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/api_worker_scheduler.rs#L187) | +| Declared at | [`nativelink-util/src/metrics.rs:908`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L908) | +| Emitted from | [`nativelink-scheduler/src/api_worker_scheduler.rs:400`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/api_worker_scheduler.rs#L400) | ### `worker.keepalive.gap` -Seconds since the previous message from the same worker, recorded at each -message. The distribution `worker_timeout_s` should be judged against: a -timeout below the tail of this histogram evicts live workers. +Seconds between consecutive messages from a worker. | Property | Value | | --- | --- | | Instrument type | `Histogram` | | Field | `worker_keepalive_gap` | | Unit | `s` | -| Buckets | 1, 2, 3, 4, 5, 6, 8, 10, 15, 20, 30, 60 | | Prometheus series | `nativelink_worker_keepalive_gap` (not referenced by any shipped rule or dashboard) | | Emitted | Conditionally. scheduler process only | -| Declared at | `nativelink-util/src/metrics.rs`, `WorkerMetrics::worker_keepalive_gap` | -| Emitted from | `nativelink-scheduler/src/api_worker_scheduler.rs`, `refresh_lifetime` | +| Declared at | [`nativelink-util/src/metrics.rs:914`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L914) | +| Emitted from | [`nativelink-scheduler/src/api_worker_scheduler.rs:348`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/api_worker_scheduler.rs#L348) | + +**Histogram buckets:** 12 explicit boundaries, from +`1.0` to `60.0` `s`. Anything past the last boundary lands in the +`+Inf` bucket, so a quantile computed above it is a lower bound, not a +measurement. ### `worker.state.count` @@ -481,8 +478,8 @@ Connected workers currently paused or draining. | Unit | `{worker}` | | Prometheus series | `nativelink_worker_state_count` (not referenced by any shipped rule or dashboard) | | Emitted | Conditionally. scheduler process only | -| Declared at | [`nativelink-util/src/metrics.rs:845`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L845) | -| Emitted from | [`nativelink-scheduler/src/api_worker_scheduler.rs:221`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/api_worker_scheduler.rs#L221), [`nativelink-scheduler/src/api_worker_scheduler.rs:224`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/api_worker_scheduler.rs#L224), [`nativelink-scheduler/src/api_worker_scheduler.rs:256`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/api_worker_scheduler.rs#L256), [`nativelink-scheduler/src/api_worker_scheduler.rs:413`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/api_worker_scheduler.rs#L413) | +| Declared at | [`nativelink-util/src/metrics.rs:925`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L925) | +| Emitted from | [`nativelink-scheduler/src/api_worker_scheduler.rs:383`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/api_worker_scheduler.rs#L383), [`nativelink-scheduler/src/api_worker_scheduler.rs:436`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/api_worker_scheduler.rs#L436), [`nativelink-scheduler/src/api_worker_scheduler.rs:439`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/api_worker_scheduler.rs#L439), [`nativelink-scheduler/src/api_worker_scheduler.rs:549`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/api_worker_scheduler.rs#L549), [`nativelink-scheduler/src/api_worker_scheduler.rs:1098`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/api_worker_scheduler.rs#L1098), [`nativelink-scheduler/src/api_worker_scheduler.rs:1140`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/api_worker_scheduler.rs#L1140), [`nativelink-scheduler/src/api_worker_scheduler.rs:1276`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/api_worker_scheduler.rs#L1276) | ## gRPC server metrics @@ -491,7 +488,7 @@ Global gRPC serving metrics. One duration histogram covers rate, errors and latency: the count gives rate, the status attribute separates errors, and the buckets give latency. -Declared as `RPC_METRICS` in [`nativelink-util/src/metrics.rs:921`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L921), backed by the `RpcMetrics` struct. +Declared as `RPC_METRICS` in [`nativelink-util/src/metrics.rs:1025`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1025), backed by the `RpcMetrics` struct. Emitted by the `OtlpLayer` tower middleware on every gRPC server NativeLink starts, so a CAS, a scheduler and a worker API endpoint each report their own. @@ -506,8 +503,8 @@ Duration of inbound gRPC calls, by service, method and status. | Unit | `s` | | Prometheus series | `nativelink_rpc_server_duration` (not referenced by any shipped rule or dashboard) | | Emitted | Yes. whenever the code path runs | -| Declared at | [`nativelink-util/src/metrics.rs:925`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L925) | -| Emitted from | [`nativelink-util/src/telemetry.rs:375`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/telemetry.rs#L375) | +| Declared at | [`nativelink-util/src/metrics.rs:1029`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1029) | +| Emitted from | [`nativelink-util/src/telemetry.rs:375`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/telemetry.rs#L375) | **Histogram buckets:** 15 explicit boundaries, from `0.001` to `300.0` `s`. Anything past the last boundary lands in the @@ -518,7 +515,7 @@ measurement. Global scheduler metrics. -Declared as `SCHEDULER_METRICS` in [`nativelink-util/src/metrics.rs:974`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L974), backed by the `SchedulerOtlpMetrics` struct. +Declared as `SCHEDULER_METRICS` in [`nativelink-util/src/metrics.rs:1078`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1078), backed by the `SchedulerOtlpMetrics` struct. Emitted from the scheduler process once per matching pass, whether or not the pass assigned anything. @@ -534,8 +531,8 @@ while the queue is non-empty means matching is the bottleneck. | Unit | `s` | | Prometheus series | `nativelink_scheduler_matching_duration` (not referenced by any shipped rule or dashboard) | | Emitted | Conditionally. scheduler process only | -| Declared at | [`nativelink-util/src/metrics.rs:978`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L978) | -| Emitted from | [`nativelink-scheduler/src/simple_scheduler.rs:281`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/simple_scheduler.rs#L281) | +| Declared at | [`nativelink-util/src/metrics.rs:1082`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1082) | +| Emitted from | [`nativelink-scheduler/src/simple_scheduler.rs:449`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler.rs#L449) | **Histogram buckets:** 13 explicit boundaries, from `0.0001` to `10.0` `s`. Anything past the last boundary lands in the @@ -553,14 +550,130 @@ Matching passes, by result. | Unit | `{pass}` | | Prometheus series | `nativelink_scheduler_matching_passes` (not referenced by any shipped rule or dashboard) | | Emitted | Conditionally. scheduler process only | -| Declared at | [`nativelink-util/src/metrics.rs:987`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L987) | -| Emitted from | [`nativelink-scheduler/src/simple_scheduler.rs:281`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/simple_scheduler.rs#L281) | +| Declared at | [`nativelink-util/src/metrics.rs:1091`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1091) | +| Emitted from | [`nativelink-scheduler/src/simple_scheduler.rs:449`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler.rs#L449) | + +### `scheduler.unsatisfiable.queued` + +Queued actions seen in the last matching pass that no worker +connected to this scheduler could run even when idle. Per scheduler +instance, judged against that instance's own workers. + +| Property | Value | +| --- | --- | +| Instrument type | `Gauge` | +| Field | `unsatisfiable_queued` | +| Unit | `{action}` | +| Prometheus series | `nativelink_scheduler_unsatisfiable_queued` (not referenced by any shipped rule or dashboard) | +| Emitted | Conditionally. scheduler process only | +| Declared at | [`nativelink-util/src/metrics.rs:1097`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1097) | +| Emitted from | [`nativelink-scheduler/src/simple_scheduler.rs:1019`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler.rs#L1019) | + +### `scheduler.unsatisfiable.failed` + +Queued actions failed because no worker could ever run them, by +the properties that could not be satisfied. + +| Property | Value | +| --- | --- | +| Instrument type | `Counter` | +| Field | `unsatisfiable_failed` | +| Unit | `{action}` | +| Prometheus series | `nativelink_scheduler_unsatisfiable_failed` (not referenced by any shipped rule or dashboard) | +| Emitted | Conditionally. scheduler process only | +| Declared at | [`nativelink-util/src/metrics.rs:1105`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1105) | +| Emitted from | [`nativelink-scheduler/src/simple_scheduler.rs:551`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler.rs#L551) | + +### `scheduler.queue.depth` + +Queued actions the last matching pass read. + +| Property | Value | +| --- | --- | +| Instrument type | `Gauge` | +| Field | `queue_depth` | +| Unit | `{action}` | +| Prometheus series | `nativelink_scheduler_queue_depth` (not referenced by any shipped rule or dashboard) | +| Emitted | Conditionally. scheduler process only | +| Declared at | [`nativelink-util/src/metrics.rs:1111`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1111) | +| Emitted from | [`nativelink-scheduler/src/simple_scheduler.rs:1020`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler.rs#L1020) | + +### `scheduler.queue.retired` + +Queued actions retired for having no client. + +| Property | Value | +| --- | --- | +| Instrument type | `Counter` | +| Field | `queue_retired` | +| Unit | `{action}` | +| Prometheus series | `nativelink_scheduler_queue_retired` (not referenced by any shipped rule or dashboard) | +| Emitted | Conditionally. scheduler process only | +| Declared at | [`nativelink-util/src/metrics.rs:1117`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1117) | +| Emitted from | [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:597`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L597) | + +### `scheduler.matching.parked_dispatched` + +Actions a pass parked for want of room and dispatched when room +opened before the pass ended, ahead of the newer actions it was reading. + +| Property | Value | +| --- | --- | +| Instrument type | `Counter` | +| Field | `parked_dispatched` | +| Unit | `{action}` | +| Prometheus series | `nativelink_scheduler_matching_parked_dispatched` (not referenced by any shipped rule or dashboard) | +| Emitted | Conditionally. scheduler process only | +| Declared at | [`nativelink-util/src/metrics.rs:1123`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1123) | +| Emitted from | [`nativelink-scheduler/src/simple_scheduler.rs:996`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler.rs#L996) | + +### `scheduler.awaited_action.orphans` + +Listed operations whose record was gone on read, by site. + +| Property | Value | +| --- | --- | +| Instrument type | `Counter` | +| Field | `awaited_action_orphans` | +| Unit | `{operation}` | +| Prometheus series | `nativelink_scheduler_awaited_action_orphans` (not referenced by any shipped rule or dashboard) | +| Emitted | Conditionally. scheduler process only | +| Declared at | [`nativelink-util/src/metrics.rs:1131`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1131) | +| Emitted from | [`nativelink-scheduler/src/admin.rs:107`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/admin.rs#L107), [`nativelink-scheduler/src/simple_scheduler.rs:591`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler.rs#L591), [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:545`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L545), [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:1467`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L1467), [`nativelink-scheduler/src/simple_scheduler_state_manager.rs:1490`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler_state_manager.rs#L1490) | + +### `scheduler.sweep.failures` + +Abandoned-queue sweeps that failed before finishing. + +| Property | Value | +| --- | --- | +| Instrument type | `Counter` | +| Field | `sweep_failures` | +| Unit | `{sweep}` | +| Prometheus series | `nativelink_scheduler_sweep_failures` (not referenced by any shipped rule or dashboard) | +| Emitted | Conditionally. scheduler process only | +| Declared at | [`nativelink-util/src/metrics.rs:1139`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1139) | +| Emitted from | [`nativelink-scheduler/src/simple_scheduler.rs:1241`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler.rs#L1241) | + +### `scheduler.hints.resolutions` + +Actions sized by the historical resource scheduler, by source. + +| Property | Value | +| --- | --- | +| Instrument type | `Counter` | +| Field | `hint_resolutions` | +| Unit | `{action}` | +| Prometheus series | `nativelink_scheduler_hints_resolutions` (not referenced by any shipped rule or dashboard) | +| Emitted | Conditionally. scheduler process only | +| Declared at | [`nativelink-util/src/metrics.rs:1145`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1145) | +| Emitted from | [`nativelink-scheduler/src/historical_resource_scheduler.rs:513`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/historical_resource_scheduler.rs#L513), [`nativelink-scheduler/src/historical_resource_scheduler.rs:517`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/historical_resource_scheduler.rs#L517), [`nativelink-scheduler/src/historical_resource_scheduler.rs:520`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/historical_resource_scheduler.rs#L520) | ## Store tier metrics Global tiered-store metrics. -Declared as `STORE_TIER_METRICS` in [`nativelink-util/src/metrics.rs:1023`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L1023), backed by the `StoreTierMetrics` struct. +Declared as `STORE_TIER_METRICS` in [`nativelink-util/src/metrics.rs:1264`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1264), backed by the `StoreTierMetrics` struct. Emitted from the `fast_slow` store wrapper only. A deployment without a `fast_slow` store emits none of these. @@ -575,8 +688,8 @@ Reads served per tier, by result. | Unit | `{operation}` | | Prometheus series | `nativelink_store_tier_operations` (not referenced by any shipped rule or dashboard) | | Emitted | Conditionally. only through a `fast_slow` store | -| Declared at | [`nativelink-util/src/metrics.rs:1027`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L1027) | -| Emitted from | [`nativelink-store/src/fast_slow_store.rs:340`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/fast_slow_store.rs#L340), [`nativelink-store/src/fast_slow_store.rs:807`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/fast_slow_store.rs#L807), [`nativelink-store/src/fast_slow_store.rs:823`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/fast_slow_store.rs#L823), [`nativelink-store/src/fast_slow_store.rs:833`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/fast_slow_store.rs#L833), [`nativelink-store/src/fast_slow_store.rs:858`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/fast_slow_store.rs#L858), [`nativelink-store/src/fast_slow_store.rs:880`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/fast_slow_store.rs#L880) | +| Declared at | [`nativelink-util/src/metrics.rs:1268`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1268) | +| Emitted from | [`nativelink-store/src/fast_slow_store.rs:340`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/fast_slow_store.rs#L340), [`nativelink-store/src/fast_slow_store.rs:807`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/fast_slow_store.rs#L807), [`nativelink-store/src/fast_slow_store.rs:823`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/fast_slow_store.rs#L823), [`nativelink-store/src/fast_slow_store.rs:833`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/fast_slow_store.rs#L833), [`nativelink-store/src/fast_slow_store.rs:858`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/fast_slow_store.rs#L858), [`nativelink-store/src/fast_slow_store.rs:880`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/fast_slow_store.rs#L880) | ### `store.tier.io` @@ -589,14 +702,14 @@ Bytes moved per tier and direction. | Unit | `By` | | Prometheus series | `nativelink_store_tier_io` (not referenced by any shipped rule or dashboard) | | Emitted | Conditionally. only through a `fast_slow` store | -| Declared at | [`nativelink-util/src/metrics.rs:1033`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L1033) | -| Emitted from | [`nativelink-store/src/fast_slow_store.rs:349`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/fast_slow_store.rs#L349), [`nativelink-store/src/fast_slow_store.rs:824`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/fast_slow_store.rs#L824), [`nativelink-store/src/fast_slow_store.rs:859`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/fast_slow_store.rs#L859), [`nativelink-store/src/fast_slow_store.rs:881`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/fast_slow_store.rs#L881) | +| Declared at | [`nativelink-util/src/metrics.rs:1274`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1274) | +| Emitted from | [`nativelink-store/src/fast_slow_store.rs:349`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/fast_slow_store.rs#L349), [`nativelink-store/src/fast_slow_store.rs:824`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/fast_slow_store.rs#L824), [`nativelink-store/src/fast_slow_store.rs:859`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/fast_slow_store.rs#L859), [`nativelink-store/src/fast_slow_store.rs:881`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/fast_slow_store.rs#L881) | ## Health check metrics Global health-check metrics. -Declared as `HEALTH_METRICS` in [`nativelink-util/src/metrics.rs:1074`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L1074), backed by the `HealthMetrics` struct. +Declared as `HEALTH_METRICS` in [`nativelink-util/src/metrics.rs:1315`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1315), backed by the `HealthMetrics` struct. Emitted each time the health check endpoint runs, one observation per registered health indicator. @@ -611,14 +724,14 @@ Health check results, by namespace and status. | Unit | `{check}` | | Prometheus series | `nativelink_health_checks` (not referenced by any shipped rule or dashboard) | | Emitted | Yes. whenever the code path runs | -| Declared at | [`nativelink-util/src/metrics.rs:1078`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L1078) | -| Emitted from | [`nativelink-util/src/health_utils.rs:248`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/health_utils.rs#L248) | +| Declared at | [`nativelink-util/src/metrics.rs:1319`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1319) | +| Emitted from | [`nativelink-util/src/health_utils.rs:248`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/health_utils.rs#L248) | ## Connection pool metrics Global connection-pool metrics. -Declared as `CONNECTION_METRICS` in [`nativelink-util/src/metrics.rs:1105`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L1105), backed by the `ConnectionMetrics` struct. +Declared as `CONNECTION_METRICS` in [`nativelink-util/src/metrics.rs:1346`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1346), backed by the `ConnectionMetrics` struct. Emitted by the gRPC connection manager and the Redis store as they hand out and re-establish connections. A deployment that uses neither emits none of these. @@ -633,8 +746,8 @@ Free slots at the moment a connection was taken. | Unit | `{connection}` | | Prometheus series | `nativelink_connection_pool_available` (not referenced by any shipped rule or dashboard) | | Emitted | Yes. whenever the code path runs | -| Declared at | [`nativelink-util/src/metrics.rs:1109`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L1109) | -| Emitted from | [`nativelink-store/src/redis_store.rs:558`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/redis_store.rs#L558), [`nativelink-util/src/connection_manager.rs:341`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/connection_manager.rs#L341), [`nativelink-util/src/connection_manager.rs:351`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/connection_manager.rs#L351) | +| Declared at | [`nativelink-util/src/metrics.rs:1350`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1350) | +| Emitted from | [`nativelink-store/src/redis_store.rs:559`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/redis_store.rs#L559), [`nativelink-util/src/connection_manager.rs:341`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/connection_manager.rs#L341), [`nativelink-util/src/connection_manager.rs:351`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/connection_manager.rs#L351) | **Histogram buckets:** 11 explicit boundaries, from `0.0` to `512.0` `{connection}`. Anything past the last boundary lands in the @@ -652,8 +765,8 @@ Acquisitions, by pool and result. | Unit | `{acquisition}` | | Prometheus series | `nativelink_connection_pool_acquisitions` (not referenced by any shipped rule or dashboard) | | Emitted | Yes. whenever the code path runs | -| Declared at | [`nativelink-util/src/metrics.rs:1118`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L1118) | -| Emitted from | [`nativelink-store/src/redis_store.rs:558`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/redis_store.rs#L558), [`nativelink-util/src/connection_manager.rs:341`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/connection_manager.rs#L341), [`nativelink-util/src/connection_manager.rs:351`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/connection_manager.rs#L351) | +| Declared at | [`nativelink-util/src/metrics.rs:1359`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1359) | +| Emitted from | [`nativelink-store/src/redis_store.rs:559`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/redis_store.rs#L559), [`nativelink-util/src/connection_manager.rs:341`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/connection_manager.rs#L341), [`nativelink-util/src/connection_manager.rs:351`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/connection_manager.rs#L351) | ### `connection.reconnects` @@ -666,8 +779,8 @@ Reconnects, by pool. | Unit | `{reconnect}` | | Prometheus series | `nativelink_connection_reconnects` (not referenced by any shipped rule or dashboard) | | Emitted | Yes. whenever the code path runs | -| Declared at | [`nativelink-util/src/metrics.rs:1124`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs#L1124) | -| Emitted from | [`nativelink-store/src/redis_store.rs:318`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/redis_store.rs#L318) | +| Declared at | [`nativelink-util/src/metrics.rs:1365`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs#L1365) | +| Emitted from | [`nativelink-store/src/redis_store.rs:319`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/redis_store.rs#L319) | ## Attributes @@ -692,6 +805,8 @@ underscores. | `execution.instance` | `execution_instance` | `EXECUTION_INSTANCE` | open (free-form string) | | `execution.priority` | `execution_priority` | `EXECUTION_PRIORITY` | open (free-form string) | | `execution.worker_id` | `execution_worker_id` | `EXECUTION_WORKER_ID` | open (free-form string) | +| `execution.outcome` | `execution_outcome` | `EXECUTION_OUTCOME` | open (free-form string) | +| `execution.action_mnemonic` | `execution_action_mnemonic` | `EXECUTION_ACTION_MNEMONIC` | open (free-form string) | | `execution.exit_code` | `execution_exit_code` | `EXECUTION_EXIT_CODE` | open (free-form string) | | `execution.action_digest` | `execution_action_digest` | `EXECUTION_ACTION_DIGEST` | open (free-form string) | @@ -715,6 +830,7 @@ underscores. | Attribute | Prometheus label | Rust constant | Values | | --- | --- | --- | --- | | `scheduler.match.result` | `scheduler_match_result` | `SCHEDULER_MATCH_RESULT` | open (free-form string) | +| `scheduler.unsatisfiable.properties` | `scheduler_unsatisfiable_properties` | `SCHEDULER_UNSATISFIABLE_PROPERTIES` | open (free-form string) | ### Store tier attributes @@ -781,10 +897,10 @@ underscores. If anything here disagrees with the binary, the source wins: -- [`nativelink-util/src/metrics.rs`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/metrics.rs): every declaration -- [`nativelink-util/src/telemetry.rs`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-util/src/telemetry.rs): the OTLP exporter and its interval -- [`nativelink-store/src/cache_metrics_store.rs`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-store/src/cache_metrics_store.rs): every cache call site -- [`nativelink-scheduler/src/simple_scheduler_state_manager.rs`](https://github.com/TraceMachina/nativelink/blob/v1.6.5/nativelink-scheduler/src/simple_scheduler_state_manager.rs): every execution call site +- [`nativelink-util/src/metrics.rs`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/metrics.rs): every declaration +- [`nativelink-util/src/telemetry.rs`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-util/src/telemetry.rs): the OTLP exporter and its interval +- [`nativelink-store/src/cache_metrics_store.rs`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-store/src/cache_metrics_store.rs): every cache call site +- [`nativelink-scheduler/src/simple_scheduler_state_manager.rs`](https://github.com/TraceMachina/nativelink/blob/v1.7.2/nativelink-scheduler/src/simple_scheduler_state_manager.rs): every execution call site ## Where to go next diff --git a/web/apps/docs/content/docs/reference/nativelink-config/main.mdx b/web/apps/docs/content/docs/reference/nativelink-config/main.mdx index 8bfc7facac7..3f8694a5a3f 100644 --- a/web/apps/docs/content/docs/reference/nativelink-config/main.mdx +++ b/web/apps/docs/content/docs/reference/nativelink-config/main.mdx @@ -5,7 +5,7 @@ full: true --- {/* AUTOGENERATED — do not edit by hand. - Source: nativelink-config @ main (75b674f16b89525a09a35be02a027188c73f64ba) + Source: nativelink-config @ main (59cbad26) Regenerate from web/: bun --filter @nativelink/docs gen:config-reference */} @@ -924,17 +924,27 @@ Configuration for `ExperimentalMongoDB` store. | `single_use` | boolean | — | false | Accept one action, upload its results, and exit the worker process. The worker advertises one execution slot and never accepts a second action, including after a scheduler disconnect. A launcher must create a fresh container and fresh writable volumes for its replacement. Use one local worker per process, with no scheduler or services other than health checks. Reusing the container or its writable volumes does not provide isolation from the previous action. | | `experimental_buck2_file_capture` | [Buck2FileCaptureConfig](#buck2filecaptureconfig) | — | — | Capture full execution-container files for Buck2 builds only. Disabled when absent. Requires `single_use`; read the capture configuration's isolation requirements before enabling it for a worker pool. | | `timeout_handled_externally` | boolean | — | false (`NativeLink` fully handles timeouts) | If timeout is handled in `entrypoint` or another wrapper script. If set to true `NativeLink` will not honor the timeout the action requested and instead will always force kill the action after `max_action_timeout` has been reached. If this is set to false, the smaller value of the action's timeout and `max_action_timeout` will be used to which `NativeLink` will kill the action. | +| `drain_on_shutdown` | boolean | — | false | On shutdown, tell the scheduler at once that this worker is draining so it stops dispatching here, then finish the running actions (up to `max_action_timeout_s`) and exit. Off, the worker finishes first and only then announces itself, while the scheduler keeps offering it work it has to refuse. Needs a scheduler that understands the drain flag (v1.7.3+); an older one removes the worker and requeues its actions as soon as the message arrives. | +| `resource_enforcement` | [ResourceEnforcementConfig](#resourceenforcementconfig) | — | — | Kill an action that grows past its memory reservation before the pod's cgroup limit does. On cgroup v2 the container cgroup carries `memory.oom.group`, so the kernel's kill takes the worker and every action on it; this one takes only the offender and fails it with `FailedPrecondition`. Linux only; accepted and ignored elsewhere. Actions run through a persistent worker are not covered, since that process outlives the action and serves others. With either axis on, the worker reads an action's platform properties from the scheduler's `StartExecute` rather than the client's request, so a reservation the scheduler placed the action by is the one enforced and exported to the environment. | | `entrypoint` | string | — | {Use the command from the job request} | The command to execute on every execution request. This will be parsed as a command + arguments (not shell). Example: "run.sh" and a job with command: "sleep 5" will result in a command like: "run.sh sleep 5". | | `experimental_precondition_script` | string | — | — | An optional script to run before every action is processed on the worker. The value should be the full path to the script to execute and will pause all actions on the worker if it returns an exit code other than 0. If not set, then the worker will never pause and will continue to accept jobs according to the scheduler configuration. This is useful, for example, if the worker should not take any more actions until there is enough resource available on the machine to handle them. | +| `precondition_timeout_ms` | integer (uint64) | — | 30000 | Milliseconds the precondition script may run before it counts as failed and the action is refused as backpressure. A script that hangs used to hold the action forever. | +| `max_captured_output_bytes` | integer (uint64) | — | 0 | Bytes of an action's stdout or stderr kept in memory. Past this the output spills to a file under the action directory and is uploaded from disk, so an action that prints gigabytes does not grow the worker by that much. 0 keeps everything in memory. | +| `orphan_sweep_interval_s` | integer (uint64) | — | 0 | Seconds between sweeps of `work_directory` for action directories no running action owns, left by a cleanup that failed. 0 disables the sweep; the directory is still purged at startup. | +| `kill_grace_ms` | integer (uint64) | — | 5000 | Milliseconds between the SIGTERM a timed-out or cancelled action receives and the SIGKILL that follows if it is still running, so a tool that handles SIGTERM can write its own cleanup. The memory reservation kill skips the grace and sends SIGKILL at once. Linux only; elsewhere the kill is immediate. 0 takes the default. | +| `set_tmpdir` | boolean | — | true | Give every action its own `TMPDIR` under its action directory, removed with it, so concurrent actions writing the same file name under `$TMPDIR` no longer collide. A `TMPDIR` in the action's own environment still wins. Code that hard-codes `/tmp` is unaffected. | +| `persistent_workers` | [PersistentWorkersConfig](#persistentworkersconfig) | — | — | Bazel persistent workers: how an action whose platform properties carry `supports-workers=1` is run and how the pool of worker processes is sized. Unset takes every default below. | | `cas_fast_slow_store` | string | Yes | — | Underlying CAS store that the worker will use to download CAS artifacts. This store must be a `FastSlowStore`. The `fast` store must be a `FileSystemStore` because it will use hardlinks when building out the files instead of copying the files. The slow store must eventually resolve to the same store the scheduler/client uses to send job requests. | | `upload_action_result` | [UploadActionResultConfig](#uploadactionresultconfig) | — | — | Configuration for uploading action results. | | `work_directory` | string | Yes | — | The directory work jobs will be executed from. This directory will be fully managed by the worker service and will be purged on startup. This directory and the directory referenced in `local_filesystem_store_ref`'s `stores::FilesystemStore::content_path` must be on the same filesystem. Hardlinks will be used when placing files that are accessible to the jobs that are sourced from `local_filesystem_store_ref`'s `content_path`. | | `platform_properties` | map of string to [WorkerProperty](#workerproperty) | Yes | — | Properties of this worker. This configuration will be sent to the scheduler and used to tell the scheduler to restrict what should be executed on this worker. | +| `capacity` | [CapacityConfig](#capacityconfig) | — | — | Derive the CPU and memory properties from the worker's own cgroup limits instead of the values above; see `CapacityConfig`. When the cgroup cannot be read (not Linux, cgroup v1, no permission) the values above stand and a warning says so; if they carry neither property the worker refuses to start rather than register unable to take any action that asks for CPU or memory. | | `additional_environment` | map of string to [EnvironmentSource](#environmentsource) | — | — | An optional mapping of environment names to set for the execution as well as those specified in the action itself. If set, will set each key as an environment variable before executing the job with the value of the environment variable being the value of the property of the action being executed of that name or the fixed value. | | `directory_cache` | [DirectoryCacheConfig](#directorycacheconfig) | — | — | Optional directory cache configuration for improving performance by caching reconstructed input directories and using hardlinks instead of rebuilding them from CAS for every action. | | `experimental_active_input_leases` | boolean | — | false (eviction behavior is unchanged) | Optional and experimental: lease every digest in an active action's input Merkle closure in the worker's locally eviction-managed CAS tiers (the `cas_fast_slow_store`'s filesystem-backed stores) until the action has finished cleanup. This prevents `Lost inputs no longer available remotely` failures caused by local fast-tier eviction while inputs are being materialized under cache pressure. | | `use_namespaces` | boolean | — | False | Whether to use namespaces to isolate the execution. This is only available on Linux. It is highly recommended as it avoids a number of issues with zombie processes and also provides additional hermeticity. If explicitly set to true and it is not supported the worker will exit with an error. | | `use_mount_namespace` | boolean | — | False | Whether to use a mount namespace to isolate the worker root. This is only available on Linux and when `use_namespaces` is true. It is highly recommended provides additional hermeticity. If explicitly set to true and it is not supported or `use_namespaces` is not set to true the worker will exit with an error. | +| `isolate_tmp` | boolean | — | True when `use_mount_namespace` is true and `/tmp` exists, | Whether to give each action a private `/tmp` inside its mount namespace: the action's own `tmp` directory (the one `TMPDIR` points at) is bound over `/tmp`, so tools that hardcode `/tmp` write there too. Concurrent actions no longer collide on predictable paths under `/tmp`, nothing leaks between actions through it, and what an action writes there is removed with the action. | ## SimpleSpec @@ -944,12 +954,17 @@ Configuration for `ExperimentalMongoDB` store. | `retain_completed_for_s` | integer (uint32) | — | 60 seconds | The amount of time to retain completed actions for in case a `WaitExecution` is called after the action has completed. | | `client_action_timeout_s` | integer (uint64) | — | 60 seconds | Mark operations as completed with error if no client has updated them within this duration. | | `enable_active_action_count_metric` | boolean | — | false | Periodically count the actions in each stage and report them as the `execution.active.count` metric. | -| `worker_timeout_s` | integer (uint64) | — | 5 seconds | Remove workers from pool once the worker has not responded in this amount of time in seconds. | +| `worker_timeout_s` | integer (uint64) | — | 10 seconds (four default keepalive intervals) | Remove workers from pool once the worker has not responded in this amount of time in seconds. Any message from the worker counts, not only keepalives. Eviction requeues everything the worker held, so keep this well above the longest pause a loaded worker can see. | | `max_action_executing_timeout_s` | integer (uint64) | — | 0 (disabled) | Maximum time (seconds) an action can stay in Executing state without any worker update before being timed out and re-queued. This applies regardless of worker keepalive status, catching cases where a worker is alive (sending keepalives) but stuck on a specific action. Set to 0 to disable (relies only on `worker_timeout_s`). | | `unacknowledged_kill_timeout_s` | integer (uint64) | — | 60 seconds | Evict a worker that has not reported back on an operation it was told to kill within this many seconds. A healthy worker acknowledges a kill in moments; one that cannot is wedged, and its keepalives would otherwise keep `worker_timeout_s` from ever firing while the dead operation holds its slot. Eviction requeues the worker's other operations. | +| `dispatch_ack_timeout_s` | integer (uint64) | — | unset (off) | Requeue a dispatched action the worker has not acknowledged within this many seconds, untried and without counting an attempt, and pause the worker. The clock is the worker's own messages: a dispatch is dated by the worker's last message before it, and the check runs on each later message, so keep this above a few keepalive intervals. Workers from v1.7.3 acknowledge every dispatch; older workers never do, so leave this unset while any are connected or their every action would be requeued. | | `unsatisfiable_action_timeout_s` | integer (uint64) | — | 0 (never fail) | Fail a queued action once no connected worker has been able to run it for this many seconds, instead of leaving it queued forever. An action counts as impossible to run when no connected worker could take it even when idle: an `exact` value no worker has, a required key no worker declares, or a `minimum` larger than any worker's total. An action that is only waiting for a busy worker is never failed by this setting, and neither is any action while no workers are connected at all. | -| `max_job_retries` | integer (uint) | — | 3 | If a job returns an internal error or times out this many times when attempting to run on a worker the scheduler will return the last error to the client. Jobs will be retried and this configuration is to help prevent one rogue job from infinitely retrying and taking up a lot of resources when the task itself is the one causing the server to go into a bad state. | +| `no_worker_action_timeout_s` | integer (uint64) | — | 0 (never fail) | Seconds a queued action waits while no worker at all is connected to this scheduler, nor, on a shared backend, to any peer, before it fails with `FailedPrecondition`. `unsatisfiable_action_timeout_s` only runs against workers that are there, so a pool that scales from zero is never failed on sight; but a provisioner that will not create a pod for a shape it cannot serve leaves such an action waiting for a fleet that never comes, and this is the bound on that. Keep it above the longest time a pool needs to bring up its first worker. | +| `max_job_retries` | integer (uint) | — | 3 | If a job returns an internal error or times out this many times the scheduler completes it with `FailedPrecondition` carrying the last error, a code clients do not retry. This is to help prevent one rogue job from infinitely retrying and taking up a lot of resources when the task itself is the one causing the server to go into a bad state. A lost worker and a memory escalation are not the action's failures and have their own budgets: `max_worker_loss_retries` and `memory_escalation.max_steps`. | +| `max_worker_loss_retries` | integer (uint) | — | 10 | How many times an action whose worker was lost (disconnected, timed out, evicted, or OOM-killed as a whole) is queued again before the scheduler completes it with `FailedPrecondition`. A lost worker is usually not the action's fault, so these do not spend `max_job_retries`; the cap only stops an action that takes a worker down every time it runs. | | `allocation_strategy` | [WorkerAllocationStrategy](#workerallocationstrategy) | — | `"least_recently_used"` | The strategy used to assign workers jobs. | +| `live_memory_veto` | string | — | unset (off) | The name of a `minimum` platform property, usually `memory_kb`, that the scheduler compares against the free memory each worker reports with its keepalive: a worker reporting less than the action asks for is skipped for that action, whatever the admission ledger says. The ledger only subtracts what actions declare; this catches the ones that declared too little. Workers older than v1.7.3 report nothing and are never vetoed. | +| `memory_escalation` | [MemoryEscalationSpec](#memoryescalationspec) | — | unset (a memory kill fails the action) | When a worker reports an action killed for memory (`KILLED_MEMORY`, by its own reservation enforcement or by the kernel), requeue the action with a larger reservation instead of failing it. Escalations have their own budget (`max_steps`) and do not spend `max_job_retries`. The reservation grows to the largest memory any connected worker advertises (or `max_kb`); the last step reserves that worker whole, memory and CPU, so the action runs alone, and only a kill there fails the action: nothing in the fleet could run it. | | `experimental_backend` | [ExperimentalSimpleSchedulerBackend](#experimentalsimpleschedulerbackend) | — | memory | The storage backend to use for the scheduler. | | `worker_match_logging_interval_s` | integer (int64) | — | `10` | Every N seconds, do logging of worker matching e.g. "worker busy", "can't find any worker" Defaults to 10s. Can be set to `-1` to disable | | `fallback_match_interval_s` | integer (int64) | — | `5` | Every N seconds, run a worker matching pass even if no task or worker change notification arrived. This is a safety net for missed notifications and for scheduler backends with eventually consistent searches (for example Redis), where an operation that was re-queued may not be visible to the search triggered by its own notification. Without this, such an operation can stay queued until an unrelated event triggers another matching pass. Defaults to 5s. Zero or any negative value disables it. | @@ -988,8 +1003,14 @@ build at the main scheduler directly though. | --- | --- | --- | --- | --- | | `hints_file` | string | Yes | — | JSON file containing historical resource hints keyed by Bazel `RequestMetadata` `target_id` and/or `action_mnemonic`. | | `refresh_interval_s` | integer (uint64) | — | 30 seconds | Reload interval for `hints_file`. Set to 0 to load once. | -| `cpu_property_name` | string | — | `cpu_count` | Platform property name used for CPU minimum values. | -| `memory_property_name` | string | — | `memory_kb` | Platform property name used for memory minimum values, expressed in KiB. | +| `cpu_property_name` | string | — | `cpu_count` | Platform property name used for CPU minimum values. The nested scheduler must declare it as a `minimum` property (as it must the memory and disk names): an undeclared property is matched as an exact string, and an action carrying a number would then match no worker. | +| `memory_property_name` | string | — | `memory_kb` | Platform property name used for memory minimum values, expressed in KiB. Declared as `minimum` on the nested scheduler, like the CPU name. | +| `disk_property_name` | string | — | `disk_kb` | Platform property name used for disk minimum values, expressed in KiB. Declared as `minimum` on the nested scheduler, like the CPU name. | +| `classes` | array of [SizeClassSpec](#sizeclassspec) | — | — | The fleet's size ladder. A hint may name a `class` instead of numbers, and the cold-start policy below picks one for actions no hint matches. Empty means raw numbers only. | +| `default_class` | string | — | — | The class an untagged action gets when no hint and no rule below matches. Takes precedence over `cold_start` when both are set. | +| `class_by_mnemonic` | map of string to string | — | — | Class by Bazel `action_mnemonic` for untagged actions; a known heavy mnemonic (`TestRunner`, `Link`) starts in the right class. | +| `class_by_timeout` | array of [TimeoutClassSpec](#timeoutclassspec) | — | — | Class by the action's timeout, the largest `min_timeout_s` at or below the timeout wins. Checked before `default_class`, after the mnemonic rule. | +| `cold_start` | [ColdStartSpec](#coldstartspec) | — | — | Reservation given to an action that no hint matches and the client left untagged. Only a property that is absent is filled; a value the client sent is kept. Without this an untagged action costs the scheduler's ledger nothing, so any number of them can land on one worker. | | `scheduler` | [SchedulerSpec](#schedulerspec) | Yes | — | The nested scheduler to use after applying resource hints. | ## ListenerConfig @@ -1153,6 +1174,7 @@ Configuration for an individual shard of the store. | --- | --- | | `"cas"` | The store is content addressable storage. | | `"ac"` | The store is an action cache. | +| `"event_sink"` | A write-only event sink using `ByteStream.Write`. Only string keys and exact upload sizes are accepted. The resource hash identifies the key, not the payload; the size describes the payload. The upstream must be an event receiver, not a CAS. Reads and existence checks are unsupported. Use this for BEP/origin event forwarding, never for cache data. | ## GrpcReadBatchingConfig @@ -1209,6 +1231,30 @@ archive is partial. See `nativelink-worker/README.md` for the control protocol. | `internal_paths` | array of string | — | — | Internal CAS/cache paths excluded from browsing. Materialized action inputs hard-linked from these caches remain visible in the action tree. | | `finalize_timeout_s` | integer (uint) | — | 120 seconds (zero selects this default) | Maximum final archive time before action cleanup proceeds. The helper records a partial capture on timeout; build success/failure is unchanged. | +## ResourceEnforcementConfig + +| Field | Type | Required | Default | Description | +| --- | --- | --- | --- | --- | +| `memory` | [MemoryEnforcement](#memoryenforcement) | — | soft | What to do with an action that exceeds its memory reservation. | +| `disk` | [DiskEnforcement](#diskenforcement) | — | `"none"` | Whether to refuse an action whose disk reservation exceeds the free space under the work directory, less what the actions already admitted reserved, before its inputs are fetched. The refusal is `ResourceExhausted`, which the scheduler requeues without counting an attempt and holds off this worker until its next keepalive. | +| `disk_property_name` | string | — | `disk_kb` | Platform property carrying the action's disk reservation in KiB. | +| `memory_property_name` | string | — | `memory_kb` | Platform property carrying the action's memory reservation in KiB, as the scheduler sends it. An action without the property is not enforced. | +| `memory_headroom_percent` | integer (uint64) | — | 20 | Percent above the reservation an action may reach before it is killed. | +| `disk_headroom_percent` | integer (uint64) | — | 20 | Percent above its disk reservation an action's own files may reach before it is killed, with `disk: soft`. | + +## PersistentWorkersConfig + +The pool of persistent worker processes (Bazel `supports-workers`). + +| Field | Type | Required | Default | Description | +| --- | --- | --- | --- | --- | +| `enabled` | boolean | — | true | Run an action whose platform properties carry `supports-workers=1` in a pooled worker process that outlives it and serves the next action with the same executable, startup arguments and environment. Off, every such action runs as a one-shot process like any other. | +| `max_workers_per_key` | integer (uint) | — | 4 | Most worker processes kept per key (executable, startup arguments, environment, protocol). An action whose key is at the cap and has no idle process waits `acquire_timeout_s` for one to come back, then runs one-shot. 0 takes the default. | +| `idle_timeout_s` | integer (uint64) | — | 300 | Seconds an idle worker process is kept before the sweeper shuts it down. 0 takes the default. | +| `max_requests_per_worker` | integer (uint64) | — | 200 | Requests a worker process serves before it is retired and replaced, which bounds what a long-lived process accumulates. 0 takes the default. | +| `shutdown_grace_ms` | integer (uint64) | — | 5000 | Milliseconds a worker process being shut down gets to exit on its own (its stdin is closed) before SIGKILL. 0 takes the default. | +| `acquire_timeout_s` | integer (uint64) | — | 30 | Seconds an action waits for a worker process to come back when its key is at `max_workers_per_key` with none idle, before it runs one-shot instead. 0 takes the default. | + ## UploadActionResultConfig | Field | Type | Required | Default | Description | @@ -1239,6 +1285,24 @@ A dynamic configuration. The string will be executed as a command **Type:** string +## CapacityConfig + +Advertise CPU and memory from what the worker can actually see. The +worker reads `cpu.max` and `memory.max` at its own cgroup v2 root (the +pod's limits on Kubernetes), takes off what it needs for itself, divides +the memory by the enforcement headroom, and sets the two properties to +the result at registration. What the scheduler packs against is then +derived from the one number that is enforced, instead of typed in twice. + +| Field | Type | Required | Default | Description | +| --- | --- | --- | --- | --- | +| `cpu_property_name` | string | — | `cpu_count` | Property that carries CPU, on the scale `cpu_unit` names. | +| `memory_property_name` | string | — | `memory_kb` | Property that carries memory, in KiB. | +| `cpu_unit` | [CpuUnit](#cpuunit) | — | cores | The scale the CPU property is advertised on. Whole cores by default, the scale every other example in this configuration uses; a fleet that advertises and requests thousandths of a core sets `millicores`. Getting this wrong is a thousandfold error in how many actions the scheduler packs onto the worker. | +| `overhead_cpu_millicores` | integer (uint64) | — | 1000 | CPU the worker keeps for itself, in thousandths of a core whatever `cpu_unit` says. | +| `overhead_memory_kb` | integer (uint64) | — | 2097152 (2 GiB) | Memory the worker keeps for itself (the process, its directory cache, output buffers), in KiB. | +| `memory_headroom_percent` | integer (uint64) | — | unset | Percent the advertised memory is reduced by, so that actions at their enforcement ceiling (reservation plus headroom) still fit the limit. Unset, it follows `resource_enforcement.memory_headroom_percent` while memory enforcement is on, and is 0 otherwise, so the two cannot drift apart by being typed twice. | + ## EnvironmentSource One of: @@ -1283,6 +1347,19 @@ workers are able to run the task. | --- | --- | | `"least_recently_used"` | Prefer workers that have been least recently used to run a job. | | `"most_recently_used"` | Prefer workers that have been most recently used to run a job. | +| `"least_loaded"` | Prefer the worker running the fewest actions, least recently used among equals. A worker that just joined has nothing running, so it takes the next action instead of waiting for every older worker to be used once more, and a burst spreads across the fleet instead of filling one worker to its concurrency cap. | +| `"best_fit"` | Prefer the worker the action fits most tightly: among the workers that can take it, the one whose remaining `minimum` properties (memory, CPU) would be smallest afterwards, as a share of what it advertises. Big actions are kept for the workers with big room and small ones fill the gaps, so a fleet of mixed sizes wastes the least. A `priority` property on the action that a worker carries with the same value prefers that worker before the fit is judged, which is what such a property is for (a zone, a cache locality). Ties go to the fewest running actions, then least recently used. | + +## MemoryEscalationSpec + +| Field | Type | Required | Default | Description | +| --- | --- | --- | --- | --- | +| `property` | string | — | `memory_kb` | The `minimum` platform property carrying the memory reservation. | +| `percent` | integer (uint64) | — | 200 | Each kill scales the reservation by this factor, in hundredths: 200 doubles it, 150 adds half. | +| `max_kb` | integer (uint64) | — | 0 | Never reserve more than this many KiB; 0 takes the largest memory any connected worker advertises. | +| `ladder_kb` | array of integer (uint64) | — | empty (scale by `percent`) | The memory values of the fleet's size classes, ascending, in KiB. When set, a kill steps the reservation to the next class above it instead of scaling by `percent`; past the top class the last step reserves the largest worker whole. The chart fills this from the same classes the `historical_resource` scheduler uses. | +| `cpu_property` | string | — | `cpu_count` | The `minimum` platform property carrying the CPU reservation. The last escalation reserves the largest worker's whole memory and, through this property, its whole CPU, so nothing shares the worker with the action. Empty leaves CPU alone. | +| `max_steps` | integer (uint64) | — | 8 | The most escalations one action gets before the scheduler completes it with `FailedPrecondition`; 0 lets the ceiling alone bound them. | ## ExperimentalSimpleSchedulerBackend @@ -1337,6 +1414,38 @@ If a property is found, then replace it with another one. **Type:** [PlatformPropertyReplacement](#platformpropertyreplacement) +## SizeClassSpec + +A named point on the fleet's size ladder. List them ascending; each +should dominate the one before on every dimension it names. A zero +leaves that dimension alone. + +| Field | Type | Required | Default | Description | +| --- | --- | --- | --- | --- | +| `name` | string | Yes | — | Free text, referenced by hints and by the cold-start policy. | +| `cpu_count` | integer (uint64) | — | `0` | Value for `cpu_property_name`, on the scale the workers advertise. | +| `memory_kb` | integer (uint64) | — | `0` | Value for `memory_property_name`, in KiB. | +| `disk_kb` | integer (uint64) | — | `0` | Value for `disk_property_name`, in KiB. | + +## TimeoutClassSpec + +A class for actions whose timeout is at least `min_timeout_s`. Bazel's +test sizes arrive as timeouts (short 60 s, moderate 300 s, long 900 s, +eternal 3600 s), so this maps them to classes without a CAS fetch. + +| Field | Type | Required | Default | Description | +| --- | --- | --- | --- | --- | +| `min_timeout_s` | integer (uint64) | Yes | — | | +| `class` | string | Yes | — | | + +## ColdStartSpec + +| Field | Type | Required | Default | Description | +| --- | --- | --- | --- | --- | +| `cpu_count` | integer (uint64) | — | `0` | Value for `cpu_property_name`, on the scale the workers advertise. 0 leaves the property alone. | +| `memory_kb` | integer (uint64) | — | `0` | Value for `memory_property_name`, in KiB. 0 leaves the property alone. | +| `disk_kb` | integer (uint64) | — | `0` | Value for `disk_property_name`, in KiB. 0 leaves the property alone. | + ## HttpListener | Field | Type | Required | Default | Description | @@ -1431,6 +1540,7 @@ If a property is found, then replace it with another one. | --- | --- | --- | --- | --- | | `path` | string | — | "/status" | Path to register the health status check. If path is "/status", and your domain is "example.com", you can reach the endpoint with: [http://example.com/status](http://example.com/status). | | `timeout_seconds` | integer (uint64) | — | 5s | Timeout on health checks. | +| `readiness_path` | string | — | "/ready" | Path of the readiness check, the stricter sibling of `path`: it answers 503 while any component is still initializing, where `path` answers 200. A worker's registration with its scheduler is such a component, so a Kubernetes readiness probe on this path turns Ready only once the worker can take work. Must differ from `path`; the same path for both is refused at startup. | ## ErrorCode @@ -1471,6 +1581,21 @@ The possible error codes that might occur on an upstream request. | `key_file` | string | — | — | Path to the private key file for client authentication. | | `use_native_roots` | boolean | — | false | If set the client will use the native roots for TLS connections. | +## MemoryEnforcement + +| Value | Description | +| --- | --- | +| `"none"` | Measure only. | +| `"soft"` | Kill the action's process group once two consecutive samples exceed the reservation plus headroom, or one sample exceeds twice it. | + +## DiskEnforcement + +| Value | Description | +| --- | --- | +| `"none"` | Fetch inputs regardless of free space. | +| `"guard"` | Refuse the action when the free space under the work directory is below its reservation. | +| `"soft"` | `guard`, and kill an action whose files under its own directory grow past its reservation plus `disk_headroom_percent`. Sampled every ten seconds by walking the action directory, counting only files the action itself wrote (inputs are hard links into the CAS and are not counted). | + ## UploadCacheResultsStrategy | Value | Description | @@ -1480,6 +1605,15 @@ The possible error codes that might occur on an upstream request. | `"everything"` | Upload all action results that complete. | | `"failures_only"` | Only upload action results that fail. | +## CpuUnit + +The scale a CPU property is advertised on. + +| Value | Description | +| --- | --- | +| `"cores"` | Whole cores, rounded down: what `cpu_count: { query_cmd: "nproc" }` advertises and what an action asking for `cpu_count=1` means. | +| `"millicores"` | Thousandths of a core, for a fleet whose actions and workers already speak that scale. | + ## ExperimentalRedisSchedulerBackend | Field | Type | Required | Default | Description | diff --git a/web/apps/docs/content/docs/remote-execution/persistent-workers.mdx b/web/apps/docs/content/docs/remote-execution/persistent-workers.mdx index 754d7204e73..bfa6aea0494 100644 --- a/web/apps/docs/content/docs/remote-execution/persistent-workers.mdx +++ b/web/apps/docs/content/docs/remote-execution/persistent-workers.mdx @@ -44,16 +44,15 @@ alive between them. ## Two things to know before you start [#two-things-to-know-before-you-start] -**Almost nothing needs configuring in NativeLink.** No field anywhere in -[the configuration model](/configuration) turns this on, tunes the pool, or -selects a protocol; `LocalWorkerConfig` has no persistent-worker section at -all, and every pool parameter is a compiled-in constant. The feature is -driven by two platform properties on each action. The one thing NativeLink -does need from you is on the scheduler: those two keys must be listed in -`supported_platform_properties`, because the scheduler rejects an action that -carries a key it does not know. If you have seen a `persistent_workers` -config block in older documentation, it did not correspond to anything in the -binary. +**Little needs configuring in NativeLink.** The pool is on by default and +the protocol is chosen by two platform properties on each action. The +worker's `persistent_workers` block (v1.7.3 and later) sizes it: how many +processes per key, how long an idle one lives, how many requests one serves, +how long an action at the cap waits for a process to come back, and +`enabled: false` to run every such action one-shot. +The one thing NativeLink does need from you is on the scheduler: those two +keys must be listed in `supported_platform_properties`, because a scheduler +older than v1.6.6 rejects an action that carries a key it does not know. **The module is licensed separately.** Most of the repository is `FSL-1.1-Apache-2.0`. Every file under `nativelink-worker/src/persistent_worker/` @@ -173,19 +172,23 @@ while everything still appears to work. ## The isolation you give up [#the-isolation-you-give-up] -This is the part to read twice. NativeLink applies real isolation to ordinary -actions and **applies none of it to persistent workers**. The persistent path -returns before any of it runs. +This is the part to read twice. A persistent worker process outlives the +actions it serves, so some of the isolation an ordinary action gets cannot +apply to it, and some of it did not apply until v1.7.3. -**No namespace sandboxing.** The PID, user, IPC, UTS and mount namespaces that -ordinary actions get are not applied. `use_namespaces` and `use_mount_namespace` -have no effect on a persistent worker. +**Process namespaces, but not the mount namespace.** With `use_namespaces` +the worker process gets its own PID, user, IPC and UTS namespaces, as an +ordinary action does, so its children are reaped with it. It never gets a +mount namespace or the private `/tmp` (`isolate_tmp`): both are per action, +and the process is not. Persistent workers see the host's `/tmp` and the +whole of the worker's filesystem, whatever `use_mount_namespace` says. -**No environment control.** Ordinary actions get a cleared environment plus -exactly what the action and config specify. A persistent worker process -inherits the NativeLink worker process's own environment, and -`environment_variables` on the command, `additional_environment` on the config, -and the side-channel environment sources are all ignored. +**The action's environment, and only that.** The process is started with a +cleared environment plus the command's `environment_variables`, the way an +ordinary action is, and that environment is part of the worker's key, so two +actions that differ in it never share a process. `additional_environment` on +the config, the side-channel sources and the per-action `TMPDIR` are not +applied: they are per action, and the process is not. **A shared, long-lived working directory.** The process's own working directory is the worker's root action directory, not the per-action one; it has to be, @@ -199,6 +202,12 @@ ever serves. use it to validate or invalidate its own cached state. Anything the tool caches across requests, it caches on trust. +**Measured, not enforced.** What the process and its children use while they +serve a request is sampled and reported as the action's resource usage, so the +sizing loop learns what the tool needs. `resource_enforcement` never kills a +persistent worker for one request's reservation: the process is shared, and a +kill would take every action queued behind it. + The consequence is a real correctness requirement, not a footnote: **a persistent worker must be hermetic by construction, because nothing is enforcing it.** State that leaks between requests produces wrong outputs that are then @@ -210,32 +219,48 @@ it. ## The pool's actual behaviour [#the-pools-actual-behaviour] -Every parameter below is a compiled-in constant. None is configurable. +The pool is configured under `persistent_workers` on the worker; every value +below is its default. + +```json5 +persistent_workers: { + enabled: true, + max_workers_per_key: 4, + idle_timeout_s: 300, + max_requests_per_worker: 200, + shutdown_grace_ms: 5000, + acquire_timeout_s: 30, +} +``` -**Four workers per key.** A fifth concurrent action for the same key does not -queue; it falls back to normal one-shot execution and logs that it did. Your -build stays correct and quietly gets slower, which is worth knowing when you are -measuring. +**Four workers per key.** A fifth concurrent action for the same key waits +up to `acquire_timeout_s` for one of the four to come back, as Bazel's own +`worker_max_instances` queue does; only if none does in that time does it +fall back to normal one-shot execution, and it logs that it did. Your build +stays correct and gets slower by a tool startup per fallback; the worker's +`persistent_worker_fallbacks` counter says how often. **Two hundred requests per worker.** After that the process is retired and the next action starts a fresh one. This is deliberate recycling of accumulated JVM state, and it means a long build will restart workers periodically by design. -**No idle timeout, in practice.** The sweeper that would retire an idle worker -exists in the code but nothing calls it. A worker lives until the request cap -retires it, until it dies on its own, or until the NativeLink worker process -exits. +**Five minutes idle.** A sweeper runs on a timer and shuts down a worker +that has served nothing for `idle_timeout_s`, so a pool of JVMs left over from +one build does not hold the worker's memory through the next. **Dead workers are noticed on the next acquire**, not proactively. A JVM that OOMs while idle is discovered when the next action asks for it, and that action starts a replacement. -**Shutdown is not graceful.** Stdin is closed, giving the tool five seconds to -exit on EOF, and then it is killed outright. No `SIGTERM` is sent. +**Shutdown is not graceful.** Stdin is closed, giving the tool `shutdown_grace_ms` +to exit on EOF, and then it is killed outright. No `SIGTERM` is sent. **One request at a time per process.** Multiplexed workers are not supported; a request carrying a non-zero request ID is rejected. +**`enabled: false` turns it off.** Every action then runs one-shot, whatever it +asks for. + ## Steps [#steps] @@ -323,10 +348,13 @@ invalidating. -More than four concurrent actions for one key fall back to one-shot execution, -and one-shot execution of a worker-wrapped tool can be slower than the plain -tool. The per-key cap is not configurable, so the levers are reducing -concurrency for that mnemonic or splitting the work across more distinct keys. +More than `max_workers_per_key` concurrent actions for one key (four by +default) wait up to `acquire_timeout_s` for a process to come back and then +fall back to one-shot execution, and one-shot execution of a worker-wrapped +tool can be slower than the plain tool. Raise the cap, or reduce concurrency +for that mnemonic. The CPU time reported for a pooled action is what the +process spent on it; its peak memory is the whole process, including what +earlier requests left resident. diff --git a/web/apps/docs/scripts/lint-snippets.mjs b/web/apps/docs/scripts/lint-snippets.mjs index 963a4f157ce..fcc410a3c8b 100644 --- a/web/apps/docs/scripts/lint-snippets.mjs +++ b/web/apps/docs/scripts/lint-snippets.mjs @@ -25,14 +25,22 @@ // // Exit code is 1 if any snippet uses a key the reference does not know. -import { readFileSync, readdirSync, statSync } from "node:fs"; +import { existsSync, readFileSync, readdirSync, statSync } from "node:fs"; import { dirname, join, relative } from "node:path"; import { fileURLToPath } from "node:url"; const here = dirname(fileURLToPath(import.meta.url)); const docsRoot = join(here, ".."); const contentDir = join(docsRoot, "content/docs"); -const REFERENCE = join(contentDir, "reference/nativelink-config/index.mdx"); +/** + * The generated references to harvest field names from: the latest release's + * page, and main's on top of it, so a field added since the last release is + * known once `gen:config-reference main` has been run. + */ +const REFERENCES = [ + join(contentDir, "reference/nativelink-config/index.mdx"), + join(contentDir, "reference/nativelink-config/main.mdx"), +]; /** Keys under these parents are operator-chosen, not config fields. */ const FREEFORM_PARENTS = new Set([ @@ -83,10 +91,17 @@ function walk(dir, out = []) { return out; } -/** Every field name and tagged-variant key the generated reference documents. */ -function harvestKnownKeys(referencePath) { - const text = readFileSync(referencePath, "utf8"); +/** Every field name and tagged-variant key the generated references document. */ +function harvestKnownKeys(referencePaths) { const known = new Set(KNOWN_MAP_KEYS); + for (const referencePath of referencePaths) { + if (!existsSync(referencePath)) continue; + harvestKnownKeysFrom(readFileSync(referencePath, "utf8"), known); + } + return known; +} + +function harvestKnownKeysFrom(text, known) { // Props-table rows: `| \`field\` | type | … |` for (const m of text.matchAll(/^\|\s*`([a-z0-9_]+)`\s*\|/gm)) known.add(m[1]); @@ -246,7 +261,7 @@ function braceBalance(src) { // -------------------------------------------------------------------------- function main() { - const known = harvestKnownKeys(REFERENCE); + const known = harvestKnownKeys(REFERENCES); const files = walk(contentDir).filter((f) => !f.includes("reference/nativelink-config/")); const problems = [];