diff --git a/docs/design/client_api_execution_modes.md b/docs/design/client_api_execution_modes.md index 71f3380b84..1809301181 100644 --- a/docs/design/client_api_execution_modes.md +++ b/docs/design/client_api_execution_modes.md @@ -2,12 +2,14 @@ ## Status -Proposal / design note. Epic: **FLARE-2698** (Client API and 3rd party integration refactoring, fix version 2.9.0). +**Status: Approved** — implementation in progress for 2.9. Epic: **FLARE-2698** (Client API and 3rd party integration refactoring, fix version 2.9.0). Revision 2 (2026-07-01): incorporates review feedback — universal session setup, owner-death and receive-side contracts, forward-path payload lifecycle, cleanup policy definition, receiver-confirmed terminal outcome, per-receiver transfer budgets, configuration surface, attach auth hardening, and a re-sequenced migration plan. Revision 2.1 (2026-07-06): adds the filters-imply-materializing-hop rule to the payload lifecycle (per-(task,direction) granularity), conversion-filter placement details (site-scope filter ordering, DXO-meta round-trip state, mixed-format cost and bf16 limit, and what selects the regime — `server_expected_format` surviving as a recipe knob), the TensorStream disposition, and a CCWF caveat cross-referencing the materializing-hop rule (a wired conversion filter moves the producer role from the subprocess to the CJ). Claims verified against code. +Revision 2.2 (2026-07-07): status set to Approved. Forward-path heartbeat exemption narrowed to the payload-materialization phase via a new TASK_PAYLOAD_READY control message (heartbeats stay authoritative while user code trains); `launch_once=False` moves to launch-scoped tokens (fresh binding material per launch, invalidated at process stop); explicit mapping between the lifecycle transfer states and the F3 `TransferOutcome` vocabulary; `ClientAPIBackendSpec` classified as internal; Open Questions annotated with their deciding PRs. + ## Background NVFlare supports the Client API so training scripts can interact with FLARE through a small surface: @@ -226,7 +228,9 @@ sequenceDiagram C->>E: task E->>T: TASK_READY (task_id, FLModel ref) T->>E: TASK_ACCEPTED - Note over T: pull task payload, then train() + Note over T: pull task payload + T->>E: TASK_PAYLOAD_READY + Note over T: train() T-->>E: LOG / HEARTBEAT T->>E: RESULT_READY (result_id, transfer_id, manifest) E->>T: RESULT_ACCEPTED (control ack, not payload done) @@ -239,7 +243,10 @@ sequenceDiagram ``` Sequence — `launch_once=False`. A fresh process per task; the HELLO handshake repeats -each round (the job-scoped launch token is reused). The key difference: the per-task +each round with fresh binding material — the CJ regenerates the launch token in the +bootstrap config it already writes per launch, and stopping the previous process +invalidates its token, so a stale process that survived termination cannot +authenticate against a later launch (the token is launch-scoped, not job-scoped). The key difference: the per-task process stop **waits for the terminal transfer outcome** before killing the process — this is what replaces the legacy `download_complete_timeout` deferred stop: @@ -250,8 +257,8 @@ sequenceDiagram participant T as Trainer loop each task C->>E: task - E->>T: write config + Popen(...) [new process] - Note over T: flare.init() (job-scoped token reused) + E->>T: write config (fresh launch token) + Popen(...) [new process] + Note over T: flare.init() (per-launch token) T->>E: HELLO (launch token) E->>T: HELLO_ACCEPTED (fresh session) E->>T: TASK_READY (task_id, FLModel ref) @@ -260,7 +267,7 @@ sequenceDiagram Note over E,T: WAIT_TRANSFER_COMPLETE (hold process until terminal) E->>C: result E->>T: SHUTDOWN - Note over E,T: stop THIS process group,
next task launches a new process + Note over E,T: stop THIS process group and invalidate its token,
next task launches a new process end ``` @@ -307,6 +314,7 @@ sequenceDiagram C->>E: task E->>T: TASK_READY (task_id, FLModel ref) T->>E: TASK_ACCEPTED + T->>E: TASK_PAYLOAD_READY (payload pulled) T-->>E: LOG / HEARTBEAT (session lease) Note over T: train() T->>E: RESULT_READY (result_id, transfer_id, manifest) @@ -410,7 +418,9 @@ subprocess behavior.) ``` CJ -> trainer : TASK_READY (task id, task name, FLModel ref, params) trainer -> CJ : TASK_ACCEPTED -...trainer pulls task payload, trains... +...trainer pulls task payload... +trainer -> CJ : TASK_PAYLOAD_READY (task id — payload materialized or lazily held, training begins) +...trainer trains (heartbeats continue autonomously)... trainer -> CJ : TASK_FAILED (task id, reason — e.g. task payload download failed) [failure path] trainer -> CJ : RESULT_READY (result_id, transfer_id, manifest) CJ -> trainer : RESULT_ACCEPTED (or RESULT_REJECTED) @@ -460,7 +470,7 @@ TASK_ACCEPTED acknowledges control receipt of the task message only — it does - In external_process the executor has two liveness signals: process exit (waitpid on the process tree it owns) and heartbeat. Process exit is authoritative for "trainer is gone"; heartbeat covers a live-but-wedged trainer. - In attach the heartbeat lease is the only liveness signal the executor has. - **Precedence rule for transfers (result path):** active data-plane transfer progress counts as session liveness. While a session has an in-flight result transfer (WAIT_PAYLOAD_ACQUIRED/WAIT_TRANSFER_COMPLETE), the session lease does not expire on missed control heartbeats alone; the transfer's own idle/progress policy (shared F3 layer) governs, and only transfer failure/timeout or an explicit abort ends the session. This resolves the otherwise-contradictory pair "heartbeat timeout invalidates the session" vs "the executor must not revoke the session before terminal payload state" — and heartbeat false-positives during multi-GB result uploads are one failure class this design removes. (The legacy stack solved this with a dedicated STREAM_PROGRESS topic feeding transfer waits, deliberately separate from peer liveness; the same separation is preserved here, just owned by the shared transfer layer.) -- **Precedence rule for the forward (task-delivery) path:** the same false-positive class exists in the *other* direction and needs its own rule, because the CJ cannot see forward bytes. When TASK_READY carries a lazy ref to a multi-GB global model, the control rank pulls those bytes **directly from the producer** (server job process); the CJ that owns the session/heartbeat lease is only a relay and observes no forward transfer (the producer's progress callback fires at the server, not the CJ). A control rank busy materializing 5 GB inside `flare.receive()` therefore has no in-flight-transfer signal to feed the rule above. Rule: **the CJ must not revoke the session on control-heartbeat timeout while a TASK_ACCEPTED is outstanding** — an accepted-but-not-yet-completed task means the trainer is legitimately working (materializing the payload and/or training), and forward materialization is bounded by the trainer-local pull policy (the shared F3 idle/timeout on the trainer→producer transfer), not by the control heartbeat. This is the forward analog of the result-path rule and closes the same 5GB-materialization false-positive that `progress_aware_streaming.md` documents on `task_payload_download`. +- **Precedence rule for the forward (task-delivery) path:** the same false-positive class exists in the *other* direction and needs its own rule, because the CJ cannot see forward bytes. When TASK_READY carries a lazy ref to a multi-GB global model, the control rank pulls those bytes **directly from the producer** (server job process); the CJ that owns the session/heartbeat lease is only a relay and observes no forward transfer (the producer's progress callback fires at the server, not the CJ). A control rank busy materializing 5 GB inside `flare.receive()` therefore has no in-flight-transfer signal to feed the rule above. Rule: **the CJ must not revoke the session on control-heartbeat timeout during the payload-materialization phase — from TASK_ACCEPTED until the trainer sends TASK_PAYLOAD_READY (or TASK_FAILED)** — because materialization is bounded by the trainer-local pull policy (the shared F3 idle/timeout on the trainer→producer transfer), not by the control heartbeat; the CJ additionally bounds the phase with a materialization deadline derived as an envelope over that pull policy, so a trainer that dies mid-pull without emitting TASK_FAILED still fails the task instead of holding the session open. The exemption ends at TASK_PAYLOAD_READY: **while user code trains, heartbeats stay authoritative** — the trainer engine's heartbeat thread runs autonomously alongside training (the same autonomy the send path guarantees above), so a healthy trainer keeps its lease through arbitrarily long training and a wedged or dead one is detected at the normal heartbeat timeout rather than at the task timeout. This matters most in attach, where the heartbeat lease is the only liveness signal. This is the forward analog of the result-path rule and closes the same 5GB-materialization false-positive that `progress_aware_streaming.md` documents on `task_payload_download` — without disabling liveness for the (much longer) training phase. #### CJ Failure (Owner Death) @@ -491,9 +501,12 @@ For Cell/FOBS-backed streaming, the Client API backend should set: - `MessageHeaderKey.MSG_ROOT_ID = transfer_id` - `MessageHeaderKey.MSG_ROOT_TTL = transfer_timeout` -- DownloadService `tx_id = transfer_id` when a single transaction is used +- DownloadService transaction ids are **attempt-scoped and never reused** (duplicate + registration raises): each attempt gets a fresh `download_tx_id`, and the manifest + carries the `transfer_id -> download_tx_id/download_ref_id` mapping. `transfer_id` + is the stable cross-attempt identity and never doubles as an F3 transaction id. -If multiple DownloadService transactions are needed, the manifest must carry the mapping from transfer_id to all download_tx_id and download_ref_id values. +The manifest must carry the mapping from transfer_id to all download_tx_id and download_ref_id values (one or many transactions alike). **Result send state machine.** The happy path runs down the left; any state can branch to a terminal failure on the right. The producer (trainer) is held alive until a terminal state is reached. @@ -512,7 +525,7 @@ DONE_CLEANUP_ALLOWED [terminal — producer may exit and free payload] ABORT at any state ────────────────────────────────► ABORTED [terminal] ``` -**Terminal states are DONE_CLEANUP_ALLOWED (success), TRANSFER_FAILED, RESULT_REJECTED, and ABORTED.** (TRANSFER_COMPLETE / TRANSFER_FAILED are the transfer-layer outcomes the machine consumes; TRANSFER_COMPLETE moves the machine to DONE_CLEANUP_ALLOWED. Earlier drafts used the two vocabularies interchangeably; they are now distinct.) RESULT_ACCEPTED is reached well before cleanup is allowed — that gap is the whole point: an accepted control message is not a completed payload transfer. +**Terminal states are DONE_CLEANUP_ALLOWED (success), TRANSFER_FAILED, RESULT_REJECTED, and ABORTED.** (TRANSFER_COMPLETE / TRANSFER_FAILED are the transfer-layer outcomes the machine consumes; TRANSFER_COMPLETE moves the machine to DONE_CLEANUP_ALLOWED. Earlier drafts used the two vocabularies interchangeably; they are now distinct.) The transfer layer itself reports a third, F3-native vocabulary: the aggregate `TransferOutcome.status` values COMPLETED / FAILED / ABORTED (`transfer_outcome.py`, reusing the TransferProgressState terminal set). The mapping at that boundary is fixed and total: TRANSFER_COMPLETE ⇔ COMPLETED; TRANSFER_FAILED ⇔ FAILED **or** ABORTED (the raw `done_status` and `reason` ride along for diagnostics, but the lifecycle machine does not branch on them). Receiver truth is applied *before* the mapping: a transaction deleted by routine cleanup after every expected receiver succeeded resolves COMPLETED, not ABORTED. RESULT_ACCEPTED is reached well before cleanup is allowed — that gap is the whole point: an accepted control message is not a completed payload transfer. Logical ownership: @@ -538,6 +551,7 @@ The producer-facing wait is then an awaitable facade over those signals: `flare. - The executor must treat duplicate RESULT_READY messages as idempotent. - If the result was already accepted, reply with the current state rather than creating a second result or second transfer. - Payload download retries happen inside the data-plane mechanism and must not create a new result_id. +- A re-offered payload transfer (transfer-level retry) is a NEW DownloadService transaction under the same transfer_id; F3 tx_ids are attempt-scoped and are never reused across attempts. - TASK_READY redelivery is idempotent by task id (see Control Protocol). **Failure and timeout rules:** @@ -576,6 +590,7 @@ The forward direction — global model from server/CJ to the trainer — is wher - **TASK_READY is control-only.** It carries the task metadata and FLModel lazy refs; TASK_ACCEPTED acknowledges the control message, before any payload materialization — the same "accepted ≠ transferred" caveat as RESULT_ACCEPTED. - **The trainer's pull is governed by the same shared transfer policy.** `flare.receive()` materializes (or lazily holds) the model by pulling through the shared F3 path; idle/progress/timeout policy is the shared layer's, and the upstream producer (CJ relay or server job process) is held to the same producer-liveness rule — its refs stay alive until the trainer's pull reaches a terminal state. - **Trainer-side download failure is explicit.** If the task payload pull fails or times out, the trainer backend sends TASK_FAILED with the task id and reason; the executor fails or retries the task at the workflow's discretion. There is no silent hang and no ambiguous half-delivered task. +- **Materialization completion is explicit.** When `flare.receive()` hands the task to user code (payload materialized, or lazily held), the trainer backend sends TASK_PAYLOAD_READY with the task id. TASK_PAYLOAD_READY or TASK_FAILED ends the payload-materialization phase — and with it the forward-path heartbeat exemption (see Heartbeat and Liveness); from that point the session lease is governed by heartbeats as normal while user code trains. - **No blind resend.** The CJ does not resend TASK_READY on a timer while the session is alive (Cell request/reply plus TASK_READY idempotency replace the Pipe resend loop that caused duplicate-delivery races). Relationship to `progress_aware_streaming.md`: that design's Phase-1 wait policies are hosted today in TaskExchanger (CJ-side forward waits) and FlareAgent (trainer-side reverse waits) over an internal Pipe topic — exactly the components this design retires. The progress-tracking substrate it defines (per-transfer progress in the shared F3 layer) is the part that survives and is what the waits here consume; the Pipe-topic delivery and TaskExchanger/FlareAgent wait owners die with the legacy stack. The forward-path contract above is the re-homed replacement: readiness comes from HELLO, delivery from Cell request/reply, and materialization waits from the shared transfer layer instead of resend suppression. @@ -733,6 +748,8 @@ silently dropped. `task_script_path`/`task_script_args` (the in_process trainer entry point), the task-name mapping, and the memory knobs are carried forward from today's InProcessClientAPIExecutor/ScriptRunner so existing jobs map without loss. +**`ClientAPIBackendSpec` is internal.** The mode backends behind the executor implement a single `ClientAPIBackendSpec` interface (in `client_api_executor.py`). Its surface is frozen early so the mode backends can be built in parallel, but it is an internal extension interface, not public API in 2.9: users select and configure modes only through the ClientAPIExecutor/ScriptRunner arguments above, and no third-party backend registration point is exposed (see the rejected general Client API Launcher under Alternatives Considered). + **No parameter converters on the executor (per FLARE-2698).** The `params_exchange_format`, `params_transfer_type`, `server_expected_format`, and `from/to_nvflare_converter_id` arguments of the legacy executors are intentionally absent. Conversion between the framework-agnostic @@ -811,7 +828,7 @@ One executor component, no pipes, no launcher, no MetricRelay: LOG messages arri ### Observability -Each session exposes its lifecycle state for operators: session state (waiting HELLO / idle / task active / waiting transfer), current task_id/result_id/transfer_id, per-receiver transfer progress (from the shared layer's progress events), and last heartbeat. These surface through the standard job stats/log channels (CJ logs at state transitions with the ids above; stats pollable via the existing cell/job info mechanisms), so "why is this producer still alive" and "which receiver is stalled" are answerable from the site without a debugger. +Each session exposes its lifecycle state for operators: session state (waiting HELLO / idle / task materializing / task active / waiting transfer), current task_id/result_id/transfer_id, per-receiver transfer progress (from the shared layer's progress events), and last heartbeat. These surface through the standard job stats/log channels (CJ logs at state transitions with the ids above; stats pollable via the existing cell/job info mechanisms), so "why is this producer still alive" and "which receiver is stalled" are answerable from the site without a debugger. ## Scenario Coverage @@ -924,7 +941,7 @@ These three are not the same moment — they fire at different workflow stages Two CCWF-specific policies still need to be set (not new transport): - **Fan-out (broadcast/CSE).** The producer stays alive until all N receivers pull; a stuck or dead receiver must not pin it forever. This is enforced with the per-receiver budgets above: each receiver stage gets a workflow-supplied expected-pull window, a receiver that exhausts its budget is marked failed for the aggregate outcome, and the workflow decides whether a partial fan-out (N-1 of N) is a usable result or a task failure. The transaction TTL is the envelope over the stage windows, plus a defined abort path. -- **Disk offload.** Two distinct artifacts must not be conflated. The **producer** serves tensors from its in-memory source (a DownloadService ref, no producer-side disk file); that source is released only on the normalized terminal outcome, which is what keeps the subprocess alive long enough. The **receiver's** offloaded copy (e.g. safetensors written by the receiving side's FOBS decode) is the receiver's own storage, on its own GC lifecycle — it is not gated on the producer's outcome. The historical CCWF fragility was the *producer/subprocess* exiting before the peer finished pulling; that is fixed by the terminal-outcome gate on the producer's source. **Hard dependency:** safely re-enabling disk offload in CCWF requires the receiver-confirmed terminal status (Payload Lifecycle refinement #1, plan PR F3-2) — re-enabling on the current *producer-served EOF* outcome would reintroduce silent truncation (a receiver whose disk write fails after its last pull is invisible to the producer). This is a hard prerequisite, not merely "requires validation." +- **Disk offload.** Two distinct artifacts must not be conflated. The **producer** serves tensors from its in-memory source (a DownloadService ref, no producer-side disk file); that source is released only on the normalized terminal outcome, which is what keeps the subprocess alive long enough. The **receiver's** offloaded copy (e.g. safetensors written by the receiving side's FOBS decode) is the receiver's own storage, on its own GC lifecycle — it is not gated on the producer's outcome. The historical CCWF fragility was the *producer/subprocess* exiting before the peer finished pulling; that is fixed by the terminal-outcome gate on the producer's source. **Hard dependency:** safely re-enabling disk offload in CCWF requires the receiver-confirmed terminal status (Payload Lifecycle refinement #1 — the receiver-confirmed completion wire change) — re-enabling on the current *producer-served EOF* outcome would reintroduce silent truncation (a receiver whose disk write fails after its last pull is invisible to the producer). This is a hard prerequisite, not merely "requires validation." ## Alternatives Considered @@ -1020,12 +1037,13 @@ Core external_process and rank contract (steps 3–4): Session and failure contracts (steps 3–4): -- external_process: a HELLO carrying a wrong/stale launch token is rejected (HELLO_REJECTED) and a stale trainer from a previous run cannot bind; attach: a bad token proof fails at HELLO_PROOF (see the attach auth tests) +- external_process: a HELLO carrying a wrong/stale launch token is rejected (HELLO_REJECTED) and a stale trainer from a previous run cannot bind (under `launch_once=False`, the previous launch's token is invalidated at process stop, so a surviving per-task process cannot bind to the next launch); attach: a bad token proof fails at HELLO_PROOF (see the attach auth tests) - HELLO with a mismatched protocol version fails fast with a clear error - duplicate TASK_READY delivery does not double-deliver the task to user code - trainer-side task payload download failure produces TASK_FAILED, not a hang - CJ kill (SIGKILL) leads to trainer process-group self-termination within the grace period; CP reaping covers a disabled trainer-side grace - heartbeat loss during an active multi-GB transfer does not revoke the session while the transfer is progressing; transfer failure/timeout does +- heartbeat loss during payload materialization (TASK_ACCEPTED outstanding, TASK_PAYLOAD_READY not yet sent) does not revoke the session; sustained heartbeat loss after TASK_PAYLOAD_READY, while user code trains, does - `flare.receive()` returns None on SHUTDOWN/job end; raises the session exception on ABORT and on session loss Payload lifecycle (steps 2, 4): @@ -1050,11 +1068,13 @@ CCWF (step 6): ### Open Questions -- Exact public argument names and migration aliases. -- Scheduler batch helper library vs. standard wrapper component. -- Exact factoring of shared code with IPCAgent/IPCExchanger. -- Compatibility flag for selecting the legacy Pipe path during migration. -- Whether partial fan-out (N-1 of N receivers succeeded) should be surfaceable to CCWF workflows as a usable result or always a task failure. +Each open question is tracked to the work item that decides it; none blocks the interface freezes already landed. + +- Exact public argument names and migration aliases — decided at the `ClientAPIExecutor` interface freeze (the frozen constructor is the decision record) and the ScriptRunner `execution_mode` wiring. +- Scheduler batch helper library vs. standard wrapper component — deferred to Migration Plan step 7 (Future Enhancements); not 2.9-blocking. +- Exact factoring of shared code with IPCAgent/IPCExchanger — decided during the trainer-engine work, where the shared session/receive machinery gets its final home. +- Compatibility flag for selecting the legacy Pipe path during migration — decided with the in_process/external_process backend wiring: whether legacy selection stays purely class-name-based (per the Compatibility section's keep-legacy story) or also gets an explicit flag. +- Whether partial fan-out (N-1 of N receivers succeeded) is surfaceable to CCWF workflows as a usable result or always a task failure — decided by the per-receiver budget work's quorum surface (`min_receivers` / `quorum_met`: `completed` stays the strict all-receivers certificate) and surfaced to workflows in the CCWF refactor. ## Migration Plan diff --git a/docs/design/client_api_execution_modes_plan.md b/docs/design/client_api_execution_modes_plan.md deleted file mode 100644 index 636d574360..0000000000 --- a/docs/design/client_api_execution_modes_plan.md +++ /dev/null @@ -1,149 +0,0 @@ -# Client API Execution Modes — 2.9 Implementation Plan - -Companion to `client_api_execution_modes.md` (design). Tracks Epic **FLARE-2698** (Client API and 3rd party integration refactoring, 2.9.0). Decomposes the design's 8-step Migration Plan into **36 PRs** with dependencies, sizes, and a release cut line. Scoped against the codebase post-2.8.0; granularity calibrated against the repo's merged-PR history (2026-07-01, see Calibration below). - -Size guide: **S** <300 changed LOC, **M** 300–800, **L** 800–2000. - -## Granularity calibration - -Measured against the last ~300 merged PRs on main: median merged PR is ~140 LOC / 4 files; 65% land at ≤300 LOC; only ~13% fall in the 300–800 band. Recent comparable programs shipped as many small PRs over weeks: **recipe API ~28 PRs, tensor-offload/streaming reliability ~20, distributed provisioning/multicloud ~18** — each roughly half this program's scope. So ~36 PRs is in line with how this repo actually ships large features, and the plan's M-heavy granularity is already on the *large* side of repo norms (p75–p85). Consolidating further would create PR shapes the repo rarely merges for core code. - -An adversarial consolidation pass was run on nine candidate merges; four survived (applied below), five were rejected. The pattern: **safe merges are same-owner/same-module dedups of work two tracks scoped twice, or two halves of one contract with no independent consumer. Bad merges drag an early foundation behind a later integration point, or weld a revert-sensitive change (interface freeze, wire-protocol version skew, security fix) to unrelated code.** Deliberate non-merges are listed in Guardrails at the end. - -## Interface freezes (do these first, review across all track owners) - -1. **Protocol vocabulary** (`nvflare/client/cell/defs.py`) — the Cell control-protocol Topics/MsgKeys/CHANNEL/version, consumed by the trainer engine, external_process, and attach. Pure constants, no I/O. (The session-token/HMAC-proof *helpers* are NOT frozen here — see the auth-model note below: they land with EP-3, where the host-trust/auth-strength decision is actually made.) -2. **ClientAPIExecutor skeleton + `ClientAPIBackendSpec`** (`nvflare/app_common/executors/client_api_executor.py`). Freeze the full V1 constructor from the design's Configuration Surface even though only in_process works initially. -3. **F3 aggregate terminal outcome** (`TransferOutcome`, additive `outcome_cb`). Reuse `TransferProgressState` terminal names; `transaction_done_cb` signature untouchable (benign-TIMEOUT callers: `hci/server/binary_transfer.py`, `app_opt/job_launcher/workspace_cell_transfer.py`). - -## Wave plan - -Tracks: **F3** payload layer · **EX** executor/ScriptRunner · **TE** trainer engine · **EP** external_process · **AT** attach · **CC** CCWF · **CT** compat/docs/tests. - -### Wave 0 — foundations (no cross-deps; start immediately, fully parallel) - -| PR | Track | Size | Notes | -|---|---|---|---| -| PR-0 Land design doc + this implementation plan in docs/design/ | all | S | **Lands first, merges fast** — the design is already approved; this just puts the reference material in-repo so every subsequent PR links it. Interface-freeze sign-offs happen on TE-1/EX-2/F3-1 themselves | -| F3-1 Aggregate all-receivers terminal transfer outcome | F3 | M | Interface freeze #3. Purely additive next to FINISHED/TIMEOUT/DELETED | -| TE-1 Protocol vocabulary (`client/cell/defs.py`) | TE | S | Interface freeze #1 — the Cell control-protocol Topics/MsgKeys/CHANNEL/version only. The auth **mechanism** is deliberately NOT frozen here: the token is 3 roles (rendezvous, anti-mixup, auth) and whether external_process needs a *secret* HMAC proof at all is a host-trust decision (single-tenant: no; multi-tenant: yes) that EP-3 makes. So the proof helpers (TokenScope, compute/verify_hello_proof, combine_nonces) land with **EP-3**, and the generic generators (generate_session_token/nonce, token_digest) go to the fuel/sec consolidation (**FLARE-3017**). The stateful SessionTokenManager is attach-only (**AT-2**) | -| EX-2 ClientAPIExecutor skeleton + backend spec + analytics-event ownership | EX | M | Interface freeze #2 | -| TE-2 Bootstrap config schema + NVFLARE_CLIENT_API_CONFIG resolution | TE | M | Additive ConfigKeys; consumes EP-1's 0600 writer. Kept separate from TE-1: touches legacy-shared `client/config.py`, different revert profile | -| EP-1 0600 permissions for Client API config files | EP | S | Fixes live exposure in today's `client_api_config.json`; standalone + backportable by design | -| EP-2 TrainerProcessRunner (process-group lifecycle + PGID records) | EP | M | SIGTERM→grace→SIGKILL; preserves SubprocessLauncher's natural-exit window; Windows path platform-guarded | -| EX-1 Export `flare.get_task_name` | EX | S | Ship-today one-liner; kept out of the churn-prone skeleton PR | - -### Wave 1 — parallel tracks behind foundations - -| PR | Track | Size | Depends | -|---|---|---|---| -| F3-2 Receiver-confirmed completion + retry-aware accounting | F3 | M | F3-1. The version-skew wire change — lands as early as possible for maximum soak; capability-flag gated, both skews interop-tested | -| F3-3 Per-(transfer, receiver) acquire/idle budgets | F3 | M | F3-1. Unconditional per-receiver activity tracking. Must also settle the quorum surface for fan-out: workflows with min_responses-style policy (k-of-N receivers suffices) need either an optional min_receivers on the transaction/facade or a documented pattern of evaluating TransferOutcome.refs against their own threshold — `completed` stays the strict all-receivers certificate either way | -| F3-4 Awaitable producer transfer facade + PAYLOAD_ACQUIRED + post-completion linger | F3 | M | F3-1 (and F3-2 for the linger's receiver-confirmed path). Must CHAIN existing DOWNLOAD_COMPLETE_CB, not replace. Folds in the former F3-5 bounded post-completion linger (a leaf with no independent consumer — only this facade's release path uses it); F3-2 still lands as its own isolated PR so the wire-skew change is not welded here | -| EX-3 in_process backend (consolidate InProcessClientAPIExecutor) | EX | L | EX-2. Behavior-parity bar: "nothing user-visible" | -| TE-3 TrainerCellSession engine (handshake, heartbeat, owner-death, trainer-side authenticated teardown) | TE | L | TE-1, TE-2. Injectable clock + kill hook; AT owner co-reviews the teardown-auth tests | -| EP-5 CP-side orphan reaping of trainer PGIDs | EP | M | EP-2. PID-reuse guard via start-time record | - -### Wave 2 — contracts and wiring - -| PR | Track | Size | Depends | -|---|---|---|---| -| TE-4 TrainerCellSession task/result contracts (receive queue + TASK_READY idempotency + TASK_FAILED; RESULT_READY flow + terminal-outcome blocking + fan-out drain/shutdown gating) | TE | L | TE-3; F3-4 via a narrow stubbed wait-protocol interface. Merged from two Ms: same class, same owner, no independent consumer of either half. Guard: split fan-out drain back out if the diff passes ~1800 LOC | -| EP-3 external_process auth + CJ-side HELLO acceptance + session-scoped message enforcement | EP | M | TE-1 (vocabulary). **Owns the auth-model decision**: a per-launch rendezvous id always; a *secret* launch token proven by one-round HMAC on multi-tenant hosts (rendezvous-only + OS isolation may suffice single-tenant). Introduces the proof helpers (TokenScope, compute/verify_hello_proof, combine_nonces) here — reused later by AT-2. Session-scoped message enforcement (accept trainer messages only from the bound session) closes the live IPCAgent any-sender gap; a P0 control, not the P2 attach track | -| EX-4 ScriptRunner `execution_mode` param + launch_external_process mapping | EX | M | EX-2/3. Convert ~22 internal recipe call sites in the same PR | -| EX-5 Converter→filter migration (FLARE-2698 bullet 2) | EX | M | EX-3, EP-4. Replace executor-owned ParamsConverters with PT/TF send+receive conversion filters at the client edge (last task-data / first task-result filter); recipes auto-wire per framework only when server format ≠ exchange format; Client API boundary passes through (RAW). Filters delegate to the existing converter classes (no logic rewrite); round-trip state (tensor shapes, excluded entries) rides DXO meta (quantizer pattern), not FLContext — removes the `_ConverterContext` stub. Keep ParamsConverter ABC deprecated + a `ParamsConverterFilter` adapter for custom converter IDs. Wiring the conversion filter makes the CJ a materializing hop (design doc: payload-lifecycle rule); numpy regime excludes bf16 — large models use pytorch end-to-end. Removes params_exchange_format/server_expected_format/converter-id from the surface (already excluded from EX-2's freeze). Transfer type FULL/DIFF stays in model_registry, decided separately | -| CC-1 CCWF transfer-declaration plumbing (receiver sets, stage windows, aux passthrough) | CC | M | F3-3. Declaration-only; absent headers preserve today's defaults — its behavior-neutrality is what de-risks the CC track | -| CT-8 Session observability (state-transition logs + StatsPoolManager view) | CT | M | EX-2; extends as backends land | - -### Wave 3 — integration point - -| PR | Track | Size | Depends | -|---|---|---|---| -| TE-5 CellClientAPI backend (flare.* on the new engine, backend selection) | TE | M | TE-4, TE-2. Opt-in only; legacy defaults untouched | -| EP-4 external_process backend for ClientAPIExecutor | EP | L | EP-2/3, EX-2/3, TE-5, F3-4. The step-4 integration point; watch for XL creep — split dispatch if needed | -| EP-6 Rank contract (torchrun multi-rank, non-control-rank fail-fast) | EP | M | EP-4, TE-5 | -| CT-1 Acceptance-test harness + core external_process suite + simulator/POC smoke | CT | L | EP-4. Absorbs the EP track's E2E validation (incl. POC-mode smoke); EP owner is co-reviewer as the track's sign-off point. tests/integration_test/fast (Blossom premerge) + xdist-safe unit subset | - -### Wave 4 — validation, workflows, attach - -| PR | Track | Size | Depends | -|---|---|---|---| -| CT-2 torchrun 2-rank rank-contract CI tests (CPU/gloo) | CT | M | CT-1, EP-4. EP owner co-reviews | -| CT-3 Owner-death (CJ-kill) + payload-lifecycle E2E tests | CT | L | CT-1, F3 track, EP-4. Covers CJ-SIGKILL self-termination + CP reaping; EP owner co-reviews | -| CC-2 Swarm onto the transfer contract (remove lazy-ref machinery) | CC | L | CC-1, F3 track, EP-4. Highest-risk PR in the program; maximally isolated revert unit; retires test_lazy_ref_local_aggr / test_msg_root_ttl | -| CC-3 Cyclic onto the transfer contract | CC | S | CC-1, EP-4. Kept out of CC-1 (would drag the behavior-neutral foundation behind EP-4) and out of CC-2 (swarm revert must not drag cyclic); shared broadcast_final_result declaration coordinates behind the CC-1 helper | -| CC-4 CSE / broadcast-best fan-out (N receivers, per-receiver budgets) | CC | L | CC-1/2, TE-4 fan-out drain | -| CC-5 Re-enable CCWF tensor disk offload (terminal-outcome-gated cleanup) | CC | M | CC-2/4. Deliberately last in track; flag experimental in 2.9 | -| AT-2 Attach session manager backend (CJ side) + attach-side session enforcement | AT | L | EX-2, TE-3, EP-3 (reuses EP-3's proof helpers). Adds the stateful `SessionTokenManager` (single-use nonce issuance, attach-window expiry, single-session, invalidation) — required because attach's token is delivered out-of-band by an untrusted starter over a possibly-remote channel, unlike external_process's localhost launch | -| AT-3 Trainer-side attach flow (bootstrap config, ad-hoc cell, HELLO proof) | AT | M | AT-2, TE-5. Connects via CP parent_url (ad-hoc listeners default-disabled) | -| CT-5 Rank-contract example updates (multi-gpu/pt, pt-ddp-docker) + SLURM batch-artifact example | CT | M | EP-4. multi-gpu/pt violates the rank contract today; qwen3-vl is the reference | - -### Wave 5 — release gate - -| PR | Track | Size | Depends | -|---|---|---|---| -| AT-4 Attach hardening (rate limit, bounded proof attempts, reconnect rotation) + E2E smoke + job-config example + delivery docs | AT | L | AT-2/3. Merged from two Ms: same owner, sequential, both inside the P2 tail so they slip together. Full negative-case matrix lives once in CT-4, not here | -| CT-4 Attach + CCWF system-level acceptance suites | CT | L | CT-1, AT track, CC track. Owns the attach negative-case matrix (replay, duplicate attach_id, spoofed teardown) — written once here, deduped from AT-4. CCWF multi-site runs in slow/ (nightly) | -| CT-6 Client API docs overhaul (client_api.rst rewrite, 3rd-party/agent docs) | CT | L | Arg names frozen (steps 3–5 merged). Written against merged code, not the design doc | -| CT-7 Legacy-stack deprecation warnings | CT | S | Coverage gate met. Warnings only in 2.9; the ScriptRunner **default flip is a separate 2.10 PR** by design (different ship dates) | - -## Program PR conventions - -Every PR in the program uses this description skeleton so reviewers always have the map: - -```markdown -## What - - -## Program context -Client API Execution Modes (2.9) — PR , Wave of the plan. -Design: docs/design/client_api_execution_modes.md § -Plan: docs/design/client_api_execution_modes_plan.md -Depends on: · Unblocks: - -## Design contracts implemented - - -## Out of scope (and where it lands instead) - -``` - -Notes: the one-paragraph "why" for the whole program lives in PR-0 and gets linked, not pasted. Interface-freeze PRs (TE-1, EX-2, F3-1) additionally name the tracks that consume the frozen surface and require sign-off from those owners before merge. - -## Critical path - -F3-1 → F3-4 → TE-4 → TE-5 → EP-4 → CT-1 → CT-3 → release gate. The TE-4 merge removed one review cycle from the path. **EP-4** remains the schedule risk to watch — it's where executor, trainer engine, auth, process runner, and the F3 facade meet; its prerequisites are deliberately extracted to keep it L. - -## 2.9 cut line (recommendation) - -- **P0 — commit for 2.9** (24 PRs): F3 track (4, after folding F3-5 into F3-4) + EX track (4) + TE track (5) + EP track (6) + CT-1/2/3/6/8. Headline: *external_process and in_process on the new Cell stack with an enforceable payload lifecycle, owner-death handling, session-scoped message enforcement, and the config-permission fix*. -- **P1 — strongly target for 2.9** (6 PRs): CC track (5) + the CCWF half of CT-4. Fully severable (SWARM keeps current behavior until CC-2 lands); ship CC-5 flagged experimental. -- **P2 — stretch for 2.9, else 2.9.x/2.10** (4 PRs): AT track (3) + the attach half of CT-4. Attach is the most self-contained step; nothing in P0/P1 depends on it (its reusable auth/proof helpers ship in P0 via EP-3, and session enforcement via EP-3, regardless). -- **CT-7 warnings** (1 PR) close 2.9. **Explicitly deferred to 2.10**: the ScriptRunner default flip and any legacy-class removal — ship 2.9 opt-in with warnings; flip after a release of field soak. - -## Effort and staffing - -36 PRs. At review-inclusive rates (S ≈ 2 days, M ≈ 1 week, L ≈ 2 weeks) that is ≈ 44 engineer-weeks total; ≈ 34 for P0+P1. With 4–5 engineers owning tracks (F3, TE, EP+EX, CC, AT+CT), full parallelism after Wave 0; calendar floor is the critical path — realistically **10–14 weeks to P0+P1 complete** plus stabilization. If 2.9 code freeze is nearer, cut P2 first, then CC-4/CC-5. - -## Guardrails — deliberate non-merges - -These pairs are tempting to consolidate and must stay separate: - -- **EP-1 (0600 fix) with anything**: live credential exposure on today's path; standalone, cherry-pickable, surgical revert. -- **EX-3 (in_process backend) + EX-4 (ScriptRunner)**: ScriptRunner is the most-used public entry point (golden exported-config tests, ~22 recipe conversions); a backend parity revert must not drag the public parameter surface. -- **Anything into CC-2 (swarm refactor)**: highest-risk PR in the program; must remain a maximally isolated revert unit. -- **F3-3 (budgets) + F3-4 (facade)**: same module, but F3-4 is on the critical path every track waits on; F3-3 serves only CC-1. -- **F3-2 (receiver-confirm wire change) must not weld to F3-4's linger/exit-timing policy**: the version-skew change wants early landing and a surgical revert; the linger now lives in F3-4 (former F3-5) but F3-2 still ships isolated ahead of it. -- **EP-3 + AT-2**: same CJ-side HELLO/session machinery, but they straddle the P0/P2 boundary — AT-2 consumes EP-3's module instead. -- **EP-2 (runner) + EP-5 (CP reaping)**: different processes and blast radii; EP-5's false-positive-reap risk needs an independent revert path. -- **CT-7 warnings + default flip**: different ship dates by design (2.9 vs 2.10). - -## Cross-cutting risks and constraints (from code scouting) - -- **CI**: GitHub premerge is ubuntu-only, CPU-only, xdist — all acceptance tests must be CPU-feasible and port-dynamic; integration tests run only via Blossom (`/build`, fast/ premerge, slow/ nightly); fork PRs get no integration signal, so protocol contracts need cheap unit-level duplicates. torchrun tests: gloo only; multi-GPU/nccl validation is manual. -- **No Windows CI**: process-group termination Windows path ships platform-guarded with skipped tests — flag as a release-note caveat. -- **Version-skew interop** (F3-2): capability-flag gating with explicit old-producer/new-receiver and new-producer/old-receiver tests; confirm sends fire-and-forget. -- **Behavior parity**: in_process consolidation and ScriptRunner wiring have golden-config tests (exported job JSON byte-compare) to protect the most-used public entry point. -- **0600 writer is single-sourced**: EP-1 owns the hardened writer; TE-2 and attach bootstrap tests consume it rather than re-implementing permission handling. -- **Pass-through registration ownership** moves from ClientAPILauncherExecutor.initialize() into the new backend — coordinate EP-4 with CC-1 to avoid double registration during migration. -- **Trainer exit hang**: FlareAgent's atexit/main-thread-join watcher for non-daemon F3 threads solves a real problem (scripts that never call flare.shutdown()); TE-3/TE-5 must carry an equivalent or trainers hang at exit. diff --git a/nvflare/client/cell/defs.py b/nvflare/client/cell/defs.py index 57ae47d64f..1bbc7cae56 100644 --- a/nvflare/client/cell/defs.py +++ b/nvflare/client/cell/defs.py @@ -13,8 +13,8 @@ # limitations under the License. """Protocol vocabulary for the Client API control protocol over Cell. -This module is interface freeze #1 of the Client API Execution Modes design -(docs/design/client_api_execution_modes.md, "Control Protocol"). It is consumed by the +This module is the frozen control-protocol vocabulary of the Client API Execution +Modes design (docs/design/client_api_execution_modes.md, "Control Protocol"). It is consumed by the trainer-side Cell engine, the external_process backend, and the attach backend. It is a pure vocabulary module: constants only, no Cell/cellnet imports, no I/O. """ @@ -50,12 +50,16 @@ class Topic: HELLO_ACCEPTED = "client_api.hello_accepted" # Distinct from ERROR (a protocol/transport error): HELLO_REJECTED is a clean, semantic # auth/handshake refusal (bad proof, wrong scope, consumed/expired nonce, single-session), - # so TE-3/AT-2 state machines can tell a recoverable auth failure from a protocol fault. + # so consumers can tell a recoverable auth failure from a protocol fault. HELLO_REJECTED = "client_api.hello_rejected" # Per task (every round) TASK_READY = "client_api.task_ready" TASK_ACCEPTED = "client_api.task_accepted" + # Sent when the trainer has materialized (or lazily bound) the task payload and hands + # it to user code. Ends the payload-materialization phase: the forward-path heartbeat + # exemption stops here and heartbeats govern the session lease while user code trains. + TASK_PAYLOAD_READY = "client_api.task_payload_ready" TASK_FAILED = "client_api.task_failed" RESULT_READY = "client_api.result_ready" RESULT_ACCEPTED = "client_api.result_accepted" diff --git a/nvflare/fuel/f3/streaming/download_service.py b/nvflare/fuel/f3/streaming/download_service.py index e7064878df..28bcb0e069 100644 --- a/nvflare/fuel/f3/streaming/download_service.py +++ b/nvflare/fuel/f3/streaming/download_service.py @@ -11,6 +11,7 @@ # 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. +import dataclasses import functools import threading import time @@ -19,6 +20,7 @@ from abc import ABC, abstractmethod from typing import Any, Callable, Optional, Tuple +from nvflare.apis.fl_constant import SystemConfigs from nvflare.apis.signal import Signal from nvflare.fuel.f3.cellnet.cell import Cell from nvflare.fuel.f3.cellnet.defs import MessageHeaderKey, ReturnCode @@ -29,11 +31,15 @@ RefOutcome, TransactionDoneStatus, TransferOutcome, + TransferOutcomeReason, compute_transfer_outcome, terminal_state_for_done_status, ) from nvflare.fuel.f3.streaming.transfer_progress import TransferProgressState +from nvflare.fuel.utils.app_config_utils import get_positive_float_var +from nvflare.fuel.utils.config_service import ConfigService from nvflare.fuel.utils.log_utils import get_obj_logger +from nvflare.fuel.utils.validation_utils import check_positive_number from nvflare.security.logging import secure_format_exception OBJ_DOWNLOADER_CHANNEL = "download_service__" @@ -94,6 +100,15 @@ Unlike with Object Streamer that the object owner pushes small objects to the recipients; with Object Downloader, each recipient pulls the data from the object owner. + +Per-receiver lifecycle on the producer side:: + + unseen --first pull--> acquired --terminal serve--> final (legacy receiver) + unseen --first pull--> acquired --terminal serve--> provisional --confirm--> final + unseen --acquire budget--> final(FAILED) + acquired/provisional --idle budget or transaction TTL--> final(FAILED) + +Final statuses feed the aggregate TransferOutcome (the receipt) at settlement. """ @@ -172,6 +187,60 @@ class _PropKey: STATE = "state" DATA = "data" STATUS = "status" + # Receiver-confirmed completion. All three keys are OPTIONAL on the wire so both + # version skews interop with legacy peers: an old receiver never sends CONFIRM_CAPABLE and + # gets today's producer-served semantics; an old producer never sends CONFIRM_EXPECTED so a + # new receiver never confirms toward it. + CONFIRM = "confirm" # receiver -> producer: terminal receiver truth (a DownloadStatus value) + CONFIRM_NONCE = "confirm_nonce" # both ways: per-serve nonce binding a confirmation to ITS serve + CONFIRM_CAPABLE = "confirm_capable" # receiver -> producer, per request: will confirm if asked + CONFIRM_EXPECTED = "confirm_expected" # producer -> receiver, per reply: confirmations consumed + + +# Per-process kill-switch for receiver-confirmed completion (read once at first use; set the +# config var / env before process start and restart to change it). The wire behavior is doubly +# gated -- per-message capability advertisement AND this switch on each side -- so a field issue +# in a mixed-version fleet is mitigated by configuration + restart without a code revert. +RECEIVER_CONFIRM_CONFIG_VAR = "streaming_receiver_confirm_enabled" +_receiver_confirm_cached = None + + +def _receiver_confirm_enabled() -> bool: + global _receiver_confirm_cached + if _receiver_confirm_cached is None: + try: + _receiver_confirm_cached = bool( + ConfigService.get_bool_var( + RECEIVER_CONFIRM_CONFIG_VAR, conf=SystemConfigs.APPLICATION_CONF, default=True + ) + ) + except Exception: + # unconfigured environments (e.g. bare unit tests) default to enabled + _receiver_confirm_cached = True + return _receiver_confirm_cached + + +# Per-(transfer, receiver) budgets. System defaults resolved from config vars; explicit +# per-transaction values win. None (unset everywhere) disables enforcement for that budget -- +# the whole-transaction timeout then remains the only backstop, exactly today's behavior. +# SIZING: the idle budget must exceed the receiver's worst-case quiet period while healthy -- +# at least its chunk-retry backoff ceiling (~60s) plus terminal store/finalization time -- +# or it manufactures failures for slow-but-healthy receivers. Budgets longer than the +# transaction timeout can never fire (warned at creation). +RECEIVER_ACQUIRE_TIMEOUT_CONFIG_VAR = "streaming_receiver_acquire_timeout" +RECEIVER_IDLE_TIMEOUT_CONFIG_VAR = "streaming_receiver_idle_timeout" + +# How long settlement waits for in-flight operations to drain. Only a hung user +# callback can exhaust it; settlement then proceeds with a warning, and the id +# stays excluded until the leaked operation exits (see _terminating_txs). +OP_DRAIN_TIMEOUT = 60.0 + + +def _resolve_receiver_budget(explicit, var_name: str): + if explicit is not None: + check_positive_number(var_name, explicit) + return float(explicit) + return get_positive_float_var(var_name, default=None) class _Ref: @@ -191,6 +260,14 @@ def __init__( self.obj = obj self.num_receivers_done = 0 self.receiver_statuses = {} + # producer-served terminal statuses awaiting the receiver's confirmation; only + # finalized (confirmed or legacy-served) statuses live in receiver_statuses + self._pending_confirms = {} + # unconditional per-receiver liveness: receiver -> last activity timestamp, + # updated on every request regardless of whether a progress_cb is configured -- so a + # live receiver can no longer mask a stalled one behind the tx-wide last_active_time + self._receiver_activity = {} + self._created_time = time.time() self._downloaded_to_all_called = False self._receiver_progress = {} self._terminal_progress_state = None @@ -200,32 +277,212 @@ def mark_active(self): self.tx.mark_active() def obj_downloaded(self, to_receiver: str, status: str): - # Status recording is guarded so terminal-outcome snapshots taken on the - # monitor thread never observe a half-updated map; user callbacks run - # outside the lock. + self._finalize_receiver(to_receiver, status) + + def _finalize_receiver( + self, to_receiver: str, status: str, require_pending: bool = False, nonce: Optional[str] = None + ) -> bool: + # Recording is guarded so outcome snapshots never observe a half-updated map; + # user callbacks run outside the lock. Dedup, pending-guard, pop, record, and + # the all-done latch are ONE critical section, so a duplicate serve can never + # resurrect a pending entry around a racing finalization. with self._progress_lock: if to_receiver in self.receiver_statuses: - return + return False + if require_pending: + pending = self._pending_confirms.get(to_receiver) + if pending is None or pending[1] != nonce: + # a legitimate confirmation always follows a provisional terminal serve on + # the CURRENT life of this ref and echoes that serve's nonce; anything else + # (unsolicited, or delayed across a ref_id reuse) must not certify -- or + # poison -- this transfer + self.tx.logger.warning( + f"dropping unsolicited/stale confirmation from {to_receiver} for ref {self.rid}" + ) + return False + self._pending_confirms.pop(to_receiver, None) self.receiver_statuses[to_receiver] = status self.num_receivers_done = len(self.receiver_statuses) assert isinstance(self.tx, _Transaction) - all_done = 0 < self.tx.num_receivers <= self.num_receivers_done and not self._downloaded_to_all_called + all_done = not self._downloaded_to_all_called and self._completion_reached_locked() if all_done: self._downloaded_to_all_called = True + # Guarded like the terminal callbacks in transaction_done: a raising user + # callback on the serving path must not lose the EOF reply for this attempt, + # and a raising downloaded_to_one must not skip downloaded_to_all (the + # _downloaded_to_all_called latch above is already set and is never retried). assert isinstance(self.obj, Downloadable) - self.obj.downloaded_to_one(to_receiver, status) + _invoke_cb_safely( + self.tx.logger, + f"downloaded_to_one of {type(self.obj)} for ref {self.rid}", + self.obj.downloaded_to_one, + to_receiver, + status, + ) if all_done: # this object is done for all receivers - self.obj.downloaded_to_all() + _invoke_cb_safely( + self.tx.logger, + f"downloaded_to_all of {type(self.obj)} for ref {self.rid}", + self.obj.downloaded_to_all, + ) + return True + + def _completion_reached_locked(self) -> bool: + # Identity-aware when the transaction declared receiver_ids: completion means every + # EXPECTED receiver is final. A status from an unexpected receiver never completes + # the ref -- otherwise (expected "b" absent + unexpected "x" present) would certify + # a delivery "b" never got. Count-based only when identities are unknown. + expected = self.tx.receiver_ids + if expected: + return all(r in self.receiver_statuses for r in expected) + return 0 < self.tx.num_receivers <= self.num_receivers_done + + def obj_served(self, to_receiver: str, status: str, expect_confirm: bool): + """Records the producer-served terminal status for a receiver. + + Legacy receivers (expect_confirm=False) finalize immediately: served EOF/ERROR is the + only truth available. Confirm-capable receivers are recorded as PROVISIONAL only -- the + receiver's confirmation (obj_confirmed) finalizes them. This is what makes accounting + retry-aware: while the record is provisional, a later serve for the same receiver + overwrites it -- a lost terminal reply healed by a retry is not stuck at the first + served status -- and the confirmation supersedes any provisional state (a receiver-side + finalization failure after the last chunk turns a served-EOF SUCCESS into a confirmed + FAILED). Once the receiver confirms, its status is final: a receiver that confirms + FAILED has given up (it confirms only on its own terminal exits). + + Returns the per-serve nonce the confirmation must echo (None when finalized + immediately or already final). The nonce binds a confirmation to THIS serve of + THIS life of the ref: a stale confirmation from a previous life of a reused + ref_id carries the wrong nonce and is dropped even if the new life has its own + pending serve for the same receiver. + """ + if not expect_confirm: + self.obj_downloaded(to_receiver, status) + return None + nonce = uuid.uuid4().hex + with self._progress_lock: + if to_receiver in self.receiver_statuses: + # already finalized -- a late duplicate serve must not resurrect a provisional + return None + self._pending_confirms[to_receiver] = (status, nonce) + return nonce + + def obj_confirmed(self, to_receiver: str, status: str, nonce: Optional[str]) -> bool: + """Records the receiver-confirmed terminal status. Receiver truth wins; first confirm is final. + + Accepted only when a provisional serve is pending for this receiver AND the + confirmation echoes that serve's nonce. The nonce is what distinguishes ref + lives: without it, a delayed confirmation from a previous life of a reused + ref_id would be accepted whenever the new life happens to have its own pending + serve for the same receiver -- certifying (or poisoning) a transfer it never saw. + """ + if status not in (DownloadStatus.SUCCESS, DownloadStatus.FAILED): + self.tx.logger.error(f"ignoring confirmation with invalid status '{status}' from {to_receiver}") + return False + accepted = self._finalize_receiver(to_receiver, status, require_pending=True, nonce=nonce) + if accepted: + # the receiver's truth is the terminal progress state for this receiver + self.emit_progress( + receiver_id=to_receiver, + state=( + TransferProgressState.COMPLETED + if status == DownloadStatus.SUCCESS + else TransferProgressState.FAILED + ), + force=True, + ) + return accepted def snapshot_receiver_statuses(self) -> dict: with self._progress_lock: return dict(self.receiver_statuses) + def snapshot_pending_confirms(self) -> dict: + with self._progress_lock: + # public shape: receiver -> provisional status (the nonce is internal) + return {r: v[0] for r, v in self._pending_confirms.items()} + + def mark_receiver_active(self, receiver: str): + now = time.time() + with self._progress_lock: + self._receiver_activity[receiver] = now + tx = self.tx + with tx._stats_lock: + # tx-level activity: budgets judge the whole transaction (see the + # _receiver_last_active field comment) + tx._acquired_receivers.add(receiver) + tx._receiver_last_active[receiver] = now + + def snapshot_receiver_activity(self) -> dict: + with self._progress_lock: + return dict(self._receiver_activity) + + def enforce_budgets( + self, now: float, acquire_timeout, idle_timeout, expected_receivers, tx_acquired=None, tx_last_active=None + ) -> list: + """Finalizes FAILED for receivers whose acquire or idle budget is exhausted. + + Candidates are the declared receiver_ids plus every receiver seen anywhere on the + transaction. Idleness is judged on TRANSACTION-level activity (last request on ANY + ref): a receiver that pulled a sibling ref and went silent has no per-ref timestamp + here, and transaction-level acquisition exempts it from the acquire budget -- judging + idle per-ref would let it escape both budgets and pin the producer to the full TTL. + + A budget failure counts toward completion (via _finalize_receiver), so the aggregate + outcome resolves on the next monitor pass. This also bounds a lost fire-and-forget + confirmation (fail-closed). + + Returns: list of (receiver, reason) that were failed on this pass. + """ + tx_acquired = tx_acquired or set() + tx_last_active = tx_last_active or {} + with self._progress_lock: + final = set(self.receiver_statuses) + candidates = set(expected_receivers or ()) | tx_acquired + failures = [] + for receiver in candidates: + if receiver in final: + continue + last_active = tx_last_active.get(receiver) + if last_active is None: + # never pulled anywhere: only the acquire budget (needs declared identities) + if acquire_timeout is not None and expected_receivers and receiver in expected_receivers: + waited = now - self._created_time + if waited > acquire_timeout: + failures.append( + ( + receiver, + f"acquire budget exhausted: no pull within {acquire_timeout}s (waited {waited:.1f}s)", + ) + ) + elif idle_timeout is not None: + idle = now - last_active + if idle > idle_timeout: + failures.append( + ( + receiver, + f"idle budget exhausted: {idle:.1f}s since last transaction activity > {idle_timeout}s", + ) + ) + enforced = [] + for receiver, reason in failures: + with self.tx._stats_lock: + if self.tx._receiver_last_active.get(receiver) != tx_last_active.get(receiver): + continue # activity advanced past the snapshot: not actually idle + # _finalize_receiver dedups and pops the pending-confirm entry itself; + # a receiver finalized meanwhile (e.g. confirmed) wins -- truth over budget + if not self._finalize_receiver(receiver, DownloadStatus.FAILED): + continue + self.tx.logger.warning(f"receiver {receiver} failed for ref {self.rid}: {reason}") + self.emit_progress(receiver_id=receiver, state=TransferProgressState.FAILED, force=True) + enforced.append((receiver, reason)) + return enforced + def emit_progress( self, *, @@ -389,6 +646,10 @@ def __init__( progress_cb: Optional[Callable] = None, progress_interval: float = 30.0, outcome_cb: Optional[Callable] = None, + receiver_ids=None, + min_receivers: Optional[int] = None, + receiver_acquire_timeout: Optional[float] = None, + receiver_idle_timeout: Optional[float] = None, ): """Constructor of the transaction object. @@ -406,6 +667,47 @@ def __init__( else: self.tid = "T" + str(uuid.uuid4()) self.timeout = timeout + self.logger = get_obj_logger(self) + + # Expected receiver identities. Optional: when provided they enable the acquire + # budget (a receiver that never issues its first pull can be failed) and, if + # num_receivers is unknown (0), supply the receiver count. + if receiver_ids: + given = len(receiver_ids) + receiver_ids = tuple(dict.fromkeys(str(r) for r in receiver_ids)) # dedup, keep order + if len(receiver_ids) != given: + # almost certainly a caller bug: it believes there are more distinct receivers + self.logger.warning( + f"tx {self.tid}: duplicate receiver_ids deduplicated ({given} -> {len(receiver_ids)})" + ) + if num_receivers and num_receivers != len(receiver_ids): + raise ValueError( + f"num_receivers ({num_receivers}) does not match receiver_ids count ({len(receiver_ids)})" + ) + num_receivers = len(receiver_ids) + self.receiver_ids = receiver_ids + else: + self.receiver_ids = None + if min_receivers is not None: + if min_receivers <= 0: + raise ValueError(f"min_receivers must be positive, got {min_receivers}") + if num_receivers and min_receivers > num_receivers: + raise ValueError(f"min_receivers ({min_receivers}) exceeds num_receivers ({num_receivers})") + self.min_receivers = min_receivers + self.receiver_acquire_timeout = _resolve_receiver_budget( + receiver_acquire_timeout, RECEIVER_ACQUIRE_TIMEOUT_CONFIG_VAR + ) + self.receiver_idle_timeout = _resolve_receiver_budget(receiver_idle_timeout, RECEIVER_IDLE_TIMEOUT_CONFIG_VAR) + for name, budget in ( + ("receiver_acquire_timeout", self.receiver_acquire_timeout), + ("receiver_idle_timeout", self.receiver_idle_timeout), + ): + if budget is not None and budget >= timeout: + # the whole-transaction clock fires first: this budget is dead config + self.logger.warning( + f"tx {self.tid}: {name}={budget}s >= transaction timeout={timeout}s -- " + f"the budget can never fire and is effectively disabled" + ) self.num_receivers = num_receivers self.last_active_time = time.time() self.start_time = time.time() @@ -420,14 +722,25 @@ def __init__( self.progress_interval = float(progress_interval) self.refs = [] self._refs_lock = threading.RLock() - self.logger = get_obj_logger(self) + # The activity gate: serves, confirms, and monitor budget passes register as + # in-flight operations; settlement closes and drains the gate before the + # outcome snapshot, so late finishes are counted and nothing emits against a + # settled transaction. + self._ops_cond = threading.Condition(threading.Lock()) + self._active_ops = 0 + self._ops_closed = False + # set in transaction_done's finally. A _terminating_txs marker is releasable only + # when this is True AND no operations are in flight: "no ops" alone does not + # mean quiet -- settlement callbacks can still be emitting. + self._settlement_complete = False + # receivers that have issued at least one pull on ANY ref (monotonic; the + # transaction-level PAYLOAD_ACQUIRED fact the acquire budget and the facade read), + # and each receiver's last activity anywhere on the transaction (what the idle + # budget judges -- per-ref timestamps would let a multi-ref receiver escape) + self._acquired_receivers = set() + self._receiver_last_active = {} def mark_active(self): - """Called to update the last active time of the transaction. - - Returns: - - """ self.last_active_time = time.time() def add_total_bytes(self, byte_count: int): @@ -445,15 +758,7 @@ def add_object( obj: Downloadable, ref_id=None, ): - """Add a large object (to be downloaded) to the transaction. - - Args: - obj: the large object to be downloaded - ref_id: the ref id to be used, if specified - - Returns: - - """ + """Adds a large object (to be downloaded) to the transaction; returns its ref.""" with self._refs_lock: r = _Ref(self, obj, ref_id) self.refs.append(r) @@ -464,102 +769,203 @@ def snapshot_refs(self): with self._refs_lock: return list(self.refs) - def timed_out(self): - """Called when the transaction is timed out. + def _fail_closed_outcome(self, done_status: str) -> Optional[TransferOutcome]: + """Direct fail-closed verdict for when compute_transfer_outcome itself raised: + full transaction metadata, honest reason, empty refs (certifies nothing).""" + try: + return TransferOutcome( + tx_id=self.tid, + status=TransferProgressState.FAILED, + reason=TransferOutcomeReason.COMPUTATION_FAILED, + done_status=done_status, + num_receivers=self.num_receivers, + refs=(), + timestamp=time.time(), + min_receivers=self.min_receivers, + receiver_ids=self.receiver_ids, + ) + except Exception as ex: + self.logger.error(f"fail-closed verdict for tx {self.tid} could not be built: {ex}") + return None - Returns: + def begin_op(self) -> bool: + """Registers an in-flight operation (serve, confirm, budget pass). + Returns False if the transaction is settling or settled: the caller must treat + it as gone (same as a missing ref). Every True return must be paired with + end_op(), normally via try/finally. """ - self.transaction_done(TransactionDoneStatus.TIMEOUT) + with self._ops_cond: + if self._ops_closed: + return False + self._active_ops += 1 + return True + + def end_op(self): + with self._ops_cond: + self._active_ops -= 1 + if self._active_ops <= 0: + self._ops_cond.notify_all() + + def _drain_ops(self, timeout: float) -> bool: + """Closes the activity gate and waits for in-flight operations to finish.""" + deadline = time.time() + timeout + with self._ops_cond: + self._ops_closed = True + while self._active_ops > 0: + remaining = deadline - time.time() + if remaining <= 0: + return False + self._ops_cond.wait(remaining) + return True + + @property + def has_receiver_budgets(self) -> bool: + return self.receiver_acquire_timeout is not None or self.receiver_idle_timeout is not None + + def enforce_receiver_budgets(self, now: float): + """Evaluates per-receiver budgets across all refs; called by the monitor thread.""" + if not self.has_receiver_budgets: + return + with self._stats_lock: + tx_acquired = set(self._acquired_receivers) + tx_last_active = dict(self._receiver_last_active) + for ref in self.snapshot_refs(): + assert isinstance(ref, _Ref) + ref.enforce_budgets( + now, + self.receiver_acquire_timeout, + self.receiver_idle_timeout, + self.receiver_ids, + tx_acquired=tx_acquired, + tx_last_active=tx_last_active, + ) def is_finished(self): - """Check whether the transaction is finished (all objects are downloaded).""" + """Check whether every expected receiver has a final status on every ref. + + Identity-aware when receiver_ids were declared: statuses from unexpected + receivers never finish a ref. A transaction with no refs yet is never + finished (it terminates via timeout or deletion instead). + """ if self.num_receivers <= 0: return False - for ref in self.snapshot_refs(): + refs = self.snapshot_refs() + if not refs: + return False + for ref in refs: assert isinstance(ref, _Ref) - if ref.num_receivers_done < self.num_receivers: - return False + with ref._progress_lock: + if not ref._completion_reached_locked(): + return False return True def transaction_done(self, status: str, on_outcome=None) -> TransferOutcome: - """Called when the transaction is finished. - - Returns the aggregate TransferOutcome (see transfer_outcome.py): COMPLETED only - when every expected receiver succeeded — TransactionDoneStatus.FINISHED alone - does not certify that. The existing transaction_done_cb contract is unchanged - except that callback exceptions no longer propagate (they would kill the - monitor thread and skip source release); outcome_cb (if set) fires after it - with the computed outcome. on_outcome (used by DownloadService to record the - outcome) is invoked right after the outcome is computed — before the - potentially slow user callbacks — so pollers see the terminal outcome as soon - as the transaction stops being active. + """Settles the transaction; returns the aggregate TransferOutcome. + + COMPLETED only when every expected receiver succeeded — FINISHED alone does + not certify that. Callback exceptions never propagate (they would kill the + monitor thread and skip source release). on_outcome (records the outcome, + releasing waiters) is invoked LAST, so acting on waiter.wait() returning can + never preempt the callback chain or source release. Runs exactly once per + transaction: every terminator unlinks it from _tx_table first. """ + # drain the activity gate first: in-flight results count in the verdict + if not self._drain_ops(OP_DRAIN_TIMEOUT): + self.logger.warning( + f"tx {self.tid}: in-flight operations did not drain within {OP_DRAIN_TIMEOUT}s; settling anyway" + ) refs = self.snapshot_refs() + outcome = None + try: + # Compute the aggregate outcome from locked per-receiver snapshots before any + # user callback can observe (or mutate the world around) this transaction. A + # computation failure must not skip the cleanup emissions below, so it is + # contained here and the verdict falls back at recording time. + try: + outcome = compute_transfer_outcome( + tx_id=self.tid, + done_status=status, + num_receivers=self.num_receivers, + min_receivers=self.min_receivers, + receiver_ids=self.receiver_ids, + refs=[ + RefOutcome(ref_id=ref.rid, receiver_statuses=ref.snapshot_receiver_statuses()) for ref in refs + ], + timestamp=time.time(), + ) + except Exception as ex: + self.logger.error(f"outcome computation for tx {self.tid} raised: {secure_format_exception(ex)}") + # the fail-closed verdict is built HERE so it flows through outcome_cb + # like any verdict (the callback contract holds even on this path) + outcome = self._fail_closed_outcome(status) + + progress_state = self._progress_state_for_transaction_status(status) + if progress_state: + for ref in refs: + ref.emit_terminal_progress_for_started_receivers(progress_state) + + elapsed = time.time() - self.start_time + total_bytes = self.get_total_bytes() + size_mb = total_bytes / (1024 * 1024) + self.logger.info( + f"[server] download tx {self.tid} done: status={status} elapsed={elapsed:.2f}s " + f"size={size_mb:.1f}MB ({total_bytes:,} bytes)" + ) - # Compute (and record via on_outcome) the aggregate outcome first, from - # locked per-receiver snapshots, before any user callback runs. - outcome = compute_transfer_outcome( - tx_id=self.tid, - done_status=status, - num_receivers=self.num_receivers, - refs=[RefOutcome(ref_id=ref.rid, receiver_statuses=ref.snapshot_receiver_statuses()) for ref in refs], - timestamp=time.time(), - ) - if on_outcome: - _invoke_cb_safely(self.logger, f"outcome recording for tx {self.tid}", on_outcome, outcome) + # Snapshot base_objs BEFORE the loop so the callback receives the + # original objects. obj.transaction_done() may clear the chunk cache + # (CacheableObject.clear_cache()); the source object itself is released + # via obj.release() AFTER the callback so the callback can still + # observe it (e.g. for memory-GC notifications). + base_objs = [ref.obj.base_obj for ref in refs] - progress_state = self._progress_state_for_transaction_status(status) - if progress_state: for ref in refs: - ref.emit_terminal_progress_for_started_receivers(progress_state) - - elapsed = time.time() - self.start_time - total_bytes = self.get_total_bytes() - size_mb = total_bytes / (1024 * 1024) - self.logger.info( - f"[server] download tx {self.tid} done: status={status} elapsed={elapsed:.2f}s " - f"size={size_mb:.1f}MB ({total_bytes:,} bytes)" - ) - - # Snapshot base_objs BEFORE the loop so the callback receives the - # original objects. obj.transaction_done() may clear the chunk cache - # (CacheableObject.clear_cache()); the source object itself is released - # via obj.release() AFTER the callback so the callback can still - # observe it (e.g. for memory-GC notifications). - base_objs = [ref.obj.base_obj for ref in refs] - - for ref in refs: - obj = ref.obj - assert isinstance(obj, Downloadable) - _invoke_cb_safely( - self.logger, - f"transaction_done of {type(obj)} for tx {self.tid}", - obj.transaction_done, - self.tid, - status, - ) - - if self.transaction_done_cb: - _invoke_cb_safely( - self.logger, - f"transaction done callback for tx {self.tid}", - self.transaction_done_cb, - self.tid, - status, - base_objs, - **self.cb_kwargs, - ) + obj = ref.obj + assert isinstance(obj, Downloadable) + _invoke_cb_safely( + self.logger, + f"transaction_done of {type(obj)} for tx {self.tid}", + obj.transaction_done, + self.tid, + status, + ) - if self.outcome_cb: - _invoke_cb_safely(self.logger, f"transfer outcome callback for tx {self.tid}", self.outcome_cb, outcome) + if self.transaction_done_cb: + _invoke_cb_safely( + self.logger, + f"transaction done callback for tx {self.tid}", + self.transaction_done_cb, + self.tid, + status, + base_objs, + **self.cb_kwargs, + ) - # Release source objects after the callback so the callback can still - # reference them. This drops the last infrastructure reference to - # large objects (e.g. numpy dicts) allowing GC to reclaim them. - for ref in refs: - ref.obj.release() + if outcome is not None and self.outcome_cb: + _invoke_cb_safely(self.logger, f"transfer outcome callback for tx {self.tid}", self.outcome_cb, outcome) + except Exception as ex: + # the ceremony must never raise: a propagating exception would kill the + # monitor thread and skip the terminator's marker sync + self.logger.error(f"settlement of tx {self.tid} raised: {secure_format_exception(ex)}") + finally: + # PHASE: source release -- independent of everything above, so no failure + # in the verdict or the callbacks can pin the sources in memory; each + # release is guarded so one raising release() cannot skip its siblings + for ref in refs: + _invoke_cb_safely(self.logger, f"release of {type(ref.obj)} for tx {self.tid}", ref.obj.release) + + # PHASE: recording -- runs last and cannot raise: the fallback is direct + # dataclass construction (never the computation that may just have failed), + # so ownership is consumed and waiters resolve no matter what happened above + if outcome is None: + # last-resort belt: an exception between the computation handler and + # here left no verdict at all + outcome = self._fail_closed_outcome(status) + if on_outcome and outcome is not None: + _invoke_cb_safely(self.logger, f"outcome recording for tx {self.tid}", on_outcome, outcome) + self._settlement_complete = True return outcome @@ -594,6 +1000,68 @@ def __init__(self, tx: _Transaction): self.objects = [r.obj for r in tx.snapshot_refs()] +class TransferWaiter: + """The awaitable facade over a transaction's terminal transfer outcome. + + This is the "returns == delivered" primitive the upper layers (executor backends, + trainer engine) consume: wait() blocks -- event-driven, no polling -- until the + transaction's aggregate TransferOutcome is recorded, and the outcome is COMPLETED only + when every expected receiver succeeded (receiver-confirmed where supported, budget- and + TTL-bounded). It attaches to the outcome-recording path directly, so it composes with -- + and never replaces -- transaction_done_cb / outcome_cb / the FOBS-context + DOWNLOAD_COMPLETE_CB chain. + """ + + def __init__(self, transaction_id: str, service=None): + self.transaction_id = transaction_id + self._service = service # the DownloadService class that created this waiter + self._event = threading.Event() + self._outcome: Optional[TransferOutcome] = None + + def _resolve(self, outcome: Optional[TransferOutcome]): + self._outcome = outcome + self._event.set() + + @property + def outcome(self) -> Optional[TransferOutcome]: + """The terminal outcome, or None while the transfer is still in flight.""" + return self._outcome + + def done(self) -> bool: + return self._event.is_set() + + def acquired_receivers(self) -> set: + """Receivers that have issued at least one pull (the PAYLOAD_ACQUIRED signal, V1).""" + service = self._service if self._service is not None else DownloadService + return service.get_acquired_receivers(self.transaction_id) + + def wait(self, timeout: Optional[float] = None, linger: Optional[float] = None) -> Optional[TransferOutcome]: + """Blocks until the terminal transfer outcome is recorded. + + Args: + timeout: max seconds to wait. None waits indefinitely (callers should normally + bound this; the transaction's own TTL and per-receiver budgets bound the + producer side). + linger: optional bounded post-completion linger, applied after any FINISHED + outcome (completed or not). By termination time the sources are already + released and the refs tombstoned; what the linger preserves is the PROCESS + (and with it the tombstone window), so a receiver whose terminal EOF/ERROR + reply was lost can still retry and be replayed its recorded status before + the producer exits. Timed-out/deleted outcomes get no linger. + + Returns: the TransferOutcome; None if the wait timed out (transfer still in flight) + or the service shut down before the transaction terminated. Disambiguate the two + None cases with done(): True means terminally resolved with no outcome (nothing + will ever record for this id -- do not re-wait); False means still in flight. + """ + if not self._event.wait(timeout): + return None + outcome = self._outcome + if outcome is not None and linger and outcome.done_status == TransactionDoneStatus.FINISHED: + time.sleep(linger) + return outcome + + class DownloadService: _init_lock = threading.Lock() @@ -607,12 +1075,23 @@ class DownloadService: # time so producers can query the aggregate result after termination. Guarded by # its own lock so outcome polling never contends with the chunk-serving _tx_lock. _tx_outcomes = {} - # Current live incarnation per tx_id (registered by new_transaction), so a - # transaction that terminates concurrently with a same-id retry cannot record - # its outcome over the new incarnation. Guarded by _outcome_lock. - _tx_incarnations = {} + # The transaction entitled to record the outcome for its tx_id (registered by + # new_transaction, consumed by _record_outcome; object-identity checked, so a + # recorder whose entry is gone drops its stale outcome). Also marks the id as + # in-use while the transaction is settling. Guarded by _outcome_lock. + _outcome_owners = {} + # Waiters blocked on a transaction's terminal outcome (the awaitable facade). Guarded by + # _outcome_lock; resolved inside _record_outcome so a waiter can never miss the outcome. + _tx_waiters = {} + # Termination markers: installed by _delete_tx at unlink (every terminator) and + # releasable only when the transaction's settlement completed AND no in-flight + # operations remain -- the id stays excluded from registration through the whole + # termination window and any drain-leaked tail, even past receipt expiry (a + # leaked emission must never land under a recycled id). Released by + # _sync_termination_marker, the registration check, or the monitor reap. + # Guarded by _tx_lock. + _terminating_txs = {} _outcome_lock = threading.Lock() - _accept_outcomes = True TX_OUTCOME_TTL = 1800.0 _logger = None _tx_monitor = None @@ -629,17 +1108,6 @@ def _initialize(cls, cell: Cell): cls._tx_monitor = threading.Thread(target=cls._monitor_tx, daemon=True) cls._tx_monitor.start() - # re-enable outcome recording for the new service lifecycle. Written under - # _outcome_lock only to keep this flag's access uniform (it is read under the - # same lock in _record_outcome). This does NOT by itself close the stale-outcome - # race across shutdown/re-init: a callback that blocked on _outcome_lock during - # shutdown can still win the lock after this write and see the re-enabled flag. - # That race is closed by the live-incarnation guard in _record_outcome, not here. - # Nesting is deadlock-safe: shutdown() acquires _outcome_lock and _init_lock - # sequentially (never nested), so there is no reverse _outcome_lock -> _init_lock. - with cls._outcome_lock: - cls._accept_outcomes = True - initialized = cls._initialized_cells.get(cell) if not initialized: # register CBs @@ -661,6 +1129,10 @@ def new_transaction( progress_cb: Optional[Callable] = None, progress_interval: float = 30.0, outcome_cb: Optional[Callable] = None, + receiver_ids=None, + min_receivers: Optional[int] = None, + receiver_acquire_timeout: Optional[float] = None, + receiver_idle_timeout: Optional[float] = None, **cb_kwargs, ): cls._initialize(cell) @@ -673,15 +1145,45 @@ def new_transaction( progress_cb=progress_cb, progress_interval=progress_interval, outcome_cb=outcome_cb, + receiver_ids=receiver_ids, + min_receivers=min_receivers, + receiver_acquire_timeout=receiver_acquire_timeout, + receiver_idle_timeout=receiver_idle_timeout, ) + # tx_ids are ATTEMPT-SCOPED and single-use while known: a retry is a NEW + # transaction with a new id (the stable cross-attempt identity is the + # caller's application-level transfer id, which never enters this service), + # so nothing a dying attempt emits can be confused with a live successor. + # A duplicate id is rejected while live, settling, receipted + # (TX_OUTCOME_TTL), or leaked (see _terminating_txs). Registration is one atomic + # step -- ownership and table entry together, _tx_lock nesting _outcome_lock + # (no path nests them in the reverse order). with cls._tx_lock: - cls._tx_table[tx.tid] = tx - with cls._outcome_lock: - # a reused explicit tx_id must not surface the previous incarnation's - # outcome: purge any recorded outcome and register this incarnation as - # current so a concurrently-terminating older incarnation cannot record - cls._tx_outcomes.pop(tx.tid, None) - cls._tx_incarnations[tx.tid] = tx + leaked = cls._terminating_txs.get(tx.tid) + if leaked is not None: + with leaked._ops_cond: + still_in_flight = leaked._active_ops > 0 + if still_in_flight or not leaked._settlement_complete: + raise ValueError( + f"transaction id {tx.tid} from a previous attempt has not fully terminated " + f"(settlement still running or operations still in flight): use a new id" + ) + cls._terminating_txs.pop(tx.tid, None) + with cls._outcome_lock: + # expire an unswept receipt inline: the exclusion window is exactly + # TX_OUTCOME_TTL, not TTL plus a sweep cycle + receipt = cls._tx_outcomes.get(tx.tid) + if receipt is not None and receipt.expired(time.time(), cls.TX_OUTCOME_TTL): + cls._tx_outcomes.pop(tx.tid, None) + receipt = None + if tx.tid in cls._tx_table or tx.tid in cls._outcome_owners or receipt is not None: + raise ValueError( + f"transaction id {tx.tid} is already in use: tx_ids are attempt-scoped and " + f"must be unique -- retry with a new id (correlate attempts with an " + f"application-level transfer id instead)" + ) + cls._outcome_owners[tx.tid] = tx + cls._tx_table[tx.tid] = tx return tx.tid @classmethod @@ -715,26 +1217,34 @@ def delete_transaction(cls, transaction_id: str): if tx: tx.transaction_done(TransactionDoneStatus.DELETED, on_outcome=functools.partial(cls._record_outcome, tx=tx)) + cls._sync_termination_marker(tx) @classmethod def shutdown(cls): - """Shutdown and clean up resources. - - Returns: None - - """ + """Shuts down the service: terminates all transactions, drops all state.""" + # Table and ownership teardown are ONE atomic step (_tx_lock nesting + # _outcome_lock): a registration landing between separate critical sections + # would enter both tables and then lose its ownership, leaving a live + # transaction that can never record. with cls._tx_lock: tx_list = list(cls._tx_table.values()) for tx in tx_list: cls._delete_tx(tx) cls._finished_refs.clear() - with cls._outcome_lock: - # stop recording (a concurrent monitor iteration may be mid-termination) - # and drop recorded outcomes; recording re-enables on next _initialize - cls._accept_outcomes = False - cls._tx_outcomes.clear() - cls._tx_incarnations.clear() + with cls._outcome_lock: + # every id being torn down already carries its termination marker + # (installed by _delete_tx at unlink), so clearing the ownership and + # receipt exclusions here cannot expose an id mid-settlement. + # Clearing ownership is what stops recording: a settlement mid-flight + # on another thread finds its entry gone at _record_outcome and drops. + cls._tx_outcomes.clear() + cls._outcome_owners.clear() + # waiters resolve to None rather than hang + for waiters in cls._tx_waiters.values(): + for waiter in waiters: + waiter._resolve(None) + cls._tx_waiters.clear() with cls._init_lock: # Shutdown resets callback-registration state even when a cell is still @@ -743,10 +1253,42 @@ def shutdown(cls): for tx in tx_list: tx.transaction_done(TransactionDoneStatus.DELETED) + cls._sync_termination_marker(tx) + + @classmethod + def _sync_termination_marker(cls, tx: _Transaction): + """Called by every terminator after settlement: release the termination marker + (installed at unlink) once settlement completed and no operations remain, or + sustain it while a drain-leaked operation is still in flight.""" + with tx._ops_cond: + leaked = tx._active_ops > 0 + with cls._tx_lock: + if leaked or not tx._settlement_complete: + cls._terminating_txs[tx.tid] = tx + elif cls._terminating_txs.get(tx.tid) is tx: + cls._terminating_txs.pop(tx.tid, None) + + @classmethod + def _reap_termination_markers(cls): + with cls._tx_lock: + released = [] + for tid, tx in cls._terminating_txs.items(): + with tx._ops_cond: # _tx_lock -> _ops_cond is the established order + done = tx._settlement_complete and tx._active_ops <= 0 + if done: + released.append(tid) + for tid in released: + cls._terminating_txs.pop(tid, None) @classmethod def _delete_tx(cls, tx: _Transaction, tombstone_finished_refs: bool = False): cls._tx_table.pop(tx.tid, None) + # install the termination marker at unlink, for EVERY terminator: it covers + # the whole settlement window and the leaked-operation tail in one mechanism, + # releasable only when settlement completed AND no operations are in flight + # (_sync_termination_marker / the duplicate check / the monitor reap). Ownership and + # receipt exclusions still exist but no longer carry the window alone. + cls._terminating_txs[tx.tid] = tx # remove all refs now = time.time() if tombstone_finished_refs else None @@ -758,26 +1300,58 @@ def _delete_tx(cls, tx: _Transaction, tombstone_finished_refs: bool = False): cls._finished_refs.pop(r.rid, None) @classmethod - def _record_outcome(cls, outcome: TransferOutcome, tx: Optional[_Transaction] = None): + def get_transfer_waiter(cls, transaction_id: str) -> TransferWaiter: + """Returns an awaitable facade over the transaction's terminal outcome. + + Safe to call before or after termination: a waiter created after the outcome was + recorded resolves immediately from the outcome table. tx_ids are attempt-scoped + (never reused), so a waiter always resolves with the verdict of exactly the + attempt it named; a retrying caller acquires a new waiter for the new attempt. + """ + waiter = TransferWaiter(transaction_id, service=cls) + with cls._outcome_lock: + existing = cls._tx_outcomes.get(transaction_id) + if existing is not None: + # resolve even from an expired record: it is still the recorded truth + waiter._resolve(existing) + return waiter + if transaction_id not in cls._outcome_owners: + # unknown/expired/shut-down: nothing will ever record for this id -- + # resolve None now (waiters never hang); race-free under this lock + waiter._resolve(None) + return waiter + cls._tx_waiters.setdefault(transaction_id, []).append(waiter) + return waiter + + @classmethod + def get_acquired_receivers(cls, transaction_id: str) -> set: + """Receivers that have issued at least one pull on any ref of the transaction.""" + with cls._tx_lock: + tx = cls._tx_table.get(transaction_id) + if tx is None: + return set() + assert isinstance(tx, _Transaction) + with tx._stats_lock: + return set(tx._acquired_receivers) + + @classmethod + def _record_outcome(cls, outcome: TransferOutcome, tx: _Transaction): + # tx is required so no call site can opt out of the owner guard: + # recording is legal only for the transaction that owns the outcome slot. with cls._outcome_lock: - current = cls._tx_incarnations.get(outcome.tx_id) - if tx is not None and current is not tx: - # Record only for the live registered incarnation. current is not tx means - # this outcome belongs to a dead generation and must be dropped, covering: - # - a newer same-tx_id incarnation (a retry) registered while this - # transaction was terminating (current is the newer tx), and - # - the incarnation was already cleared -- by shutdown() (which also - # clears the outcome table) or by a prior terminal record (current is - # None). This is what closes the stale-outcome race across shutdown/ - # re-init: a callback that blocked on _outcome_lock during shutdown and - # wins the lock afterward finds no live incarnation and is dropped here, - # regardless of the _accept_outcomes flag below. - return - if current is tx: - cls._tx_incarnations.pop(outcome.tx_id, None) - if not cls._accept_outcomes: + if cls._outcome_owners.get(outcome.tx_id) is not tx: + # ownership consumed (prior record) or cleared (shutdown): stale, drop return + cls._outcome_owners.pop(outcome.tx_id, None) + # re-stamp at recording time: the TTL retention window starts when the + # receipt becomes queryable, not when the verdict was computed -- a slow + # settlement must not record a receipt that is already expired + outcome = dataclasses.replace(outcome, timestamp=time.time()) cls._tx_outcomes[outcome.tx_id] = outcome + # resolve the awaitable facade: waiters are TransferWaiter objects (no user code + # runs in _resolve), so setting them under the lock is safe and race-free + for waiter in cls._tx_waiters.pop(outcome.tx_id, ()): + waiter._resolve(outcome) @classmethod def get_transaction_outcome(cls, transaction_id: str) -> Optional[TransferOutcome]: @@ -855,9 +1429,17 @@ def _handle_download(cls, request: Message) -> Message: cls._logger.error(f"missing {_PropKey.REF_ID} in request from {requester}") return make_reply(ReturnCode.INVALID_REQUEST) + confirm_status = payload.get(_PropKey.CONFIRM) + if confirm_status is not None: + return cls._handle_confirm(rid, requester, confirm_status, payload.get(_PropKey.CONFIRM_NONCE)) + current_state = payload.get(_PropKey.STATE) with cls._tx_lock: ref = cls._ref_table.get(rid) + if ref is not None and not ref.tx.begin_op(): + # settling/settled: a serve must not start against it -- treat the ref + # as already gone (tombstone/missing handling) + ref = None if not ref: finished_status = cls._get_finished_ref_status(rid, requester) if finished_status == DownloadStatus.SUCCESS: @@ -870,63 +1452,131 @@ def _handle_download(cls, request: Message) -> Message: cls._logger.error(f"no ref found for {rid} from {requester}") return make_reply(ReturnCode.INVALID_REQUEST) - assert isinstance(ref, _Ref) - ref.mark_active() - ref.emit_progress(receiver_id=requester, state=TransferProgressState.ACTIVE) - tx = ref.tx - assert isinstance(tx, _Transaction) - - # Keep produce() outside the global transaction lock so slow chunk generation - # does not block unrelated downloads. Timeout/delete cleanup can release the - # source concurrently; if that happens, the produce exception is reported as - # a download failure for this requester. try: - rc, data, new_state = ref.obj.produce(current_state, requester) - except Exception as ex: - ref.emit_progress(receiver_id=requester, state=TransferProgressState.FAILED, force=True) - cls._logger.error( - f"Object {type(ref.obj)} encountered exception when produce: {secure_format_exception(ex)}" - ) - return make_reply(ReturnCode.PROCESS_EXCEPTION) + assert isinstance(ref, _Ref) + ref.mark_active() + ref.mark_receiver_active(requester) + ref.emit_progress(receiver_id=requester, state=TransferProgressState.ACTIVE) + tx = ref.tx + assert isinstance(tx, _Transaction) - if rc != ProduceRC.OK: - # already done - ref.obj_downloaded( - requester, status=DownloadStatus.SUCCESS if rc == ProduceRC.EOF else DownloadStatus.FAILED - ) - ref.emit_progress( - receiver_id=requester, - state=TransferProgressState.COMPLETED if rc == ProduceRC.EOF else TransferProgressState.FAILED, - force=True, - ) - return make_reply(ReturnCode.OK, body={_PropKey.STATUS: rc}) - else: - # continue — accumulate bytes for timing summary in transaction_done() - # CacheableObject returns a list of byte-chunks; FileDownloader returns raw bytes. - # Sum chunk lengths for lists (len(list) counts items, not bytes). - if data is not None: - bytes_delta = sum(len(c) for c in data) if isinstance(data, list) else len(data) - items_delta = len(data) if isinstance(data, list) else None - tx.add_total_bytes(bytes_delta) - ref.emit_progress( - receiver_id=requester, - state=TransferProgressState.ACTIVE, - bytes_delta=bytes_delta, - items_delta=items_delta, + # receiver-confirmed completion is armed only when the receiver advertised the + # capability on this request AND the local kill-switch is on + expect_confirm = bool(payload.get(_PropKey.CONFIRM_CAPABLE)) and _receiver_confirm_enabled() + + # Keep produce() outside the global transaction lock so slow chunk generation + # does not block unrelated downloads. Timeout/delete cleanup can release the + # source concurrently; if that happens, the produce exception is reported as + # a download failure for this requester. + try: + rc, data, new_state = ref.obj.produce(current_state, requester) + except Exception as ex: + ref.emit_progress(receiver_id=requester, state=TransferProgressState.FAILED, force=True) + cls._logger.error( + f"Object {type(ref.obj)} encountered exception when produce: {secure_format_exception(ex)}" + ) + return make_reply(ReturnCode.PROCESS_EXCEPTION) + + if rc != ProduceRC.OK: + # already done -- for a confirm-capable receiver this record is PROVISIONAL and the + # receiver's confirmation finalizes it; for a legacy receiver it is final (today's + # producer-served semantics) + serve_nonce = ref.obj_served( + requester, + status=DownloadStatus.SUCCESS if rc == ProduceRC.EOF else DownloadStatus.FAILED, + expect_confirm=expect_confirm, + ) + if expect_confirm and serve_nonce: + # provisional: the receiver's confirmation carries the terminal truth -- + # do not latch a terminal progress state the confirm may contradict + ref.emit_progress(receiver_id=requester, state=TransferProgressState.ACTIVE, force=True) + body = {_PropKey.STATUS: rc, _PropKey.CONFIRM_EXPECTED: True, _PropKey.CONFIRM_NONCE: serve_nonce} + else: + ref.emit_progress( + receiver_id=requester, + state=TransferProgressState.COMPLETED if rc == ProduceRC.EOF else TransferProgressState.FAILED, + force=True, + ) + body = {_PropKey.STATUS: rc} + return make_reply(ReturnCode.OK, body=body) + else: + # continue — accumulate bytes for timing summary in transaction_done() + # CacheableObject returns a list of byte-chunks; FileDownloader returns raw bytes. + # Sum chunk lengths for lists (len(list) counts items, not bytes). + if data is not None: + bytes_delta = sum(len(c) for c in data) if isinstance(data, list) else len(data) + items_delta = len(data) if isinstance(data, list) else None + tx.add_total_bytes(bytes_delta) + ref.emit_progress( + receiver_id=requester, + state=TransferProgressState.ACTIVE, + bytes_delta=bytes_delta, + items_delta=items_delta, + ) + # no CONFIRM_EXPECTED on data chunks: the receiver only consumes it from the + # terminal reply (confirms are sent only after terminal serves), so advertising + # per chunk would be dead weight on the hottest wire message + return make_reply( + ReturnCode.OK, + body={ + _PropKey.STATUS: rc, + _PropKey.STATE: new_state, + _PropKey.DATA: data, + }, ) - return make_reply( - ReturnCode.OK, - body={ - _PropKey.STATUS: rc, - _PropKey.STATE: new_state, - _PropKey.DATA: data, - }, - ) + + finally: + ref.tx.end_op() + + @classmethod + def _handle_confirm(cls, rid: str, requester: str, status: str, nonce: Optional[str]) -> Message: + with cls._tx_lock: + ref = cls._ref_table.get(rid) + if ref is not None and not ref.tx.begin_op(): + # settling/settled: the outcome snapshot is being (or was) taken -- this + # confirm can no longer influence it and must not emit against the tx + ref = None + if ref is None: + # the transaction already terminated/cleaned up: its outcome was computed from what + # was known then (fail-closed for unconfirmed receivers); a late confirm is dropped + cls._logger.debug(f"late confirmation for unknown ref {rid} from {requester} dropped") + return make_reply(ReturnCode.OK) + assert isinstance(ref, _Ref) + try: + # deliberately no unconditional mark_active/mark_receiver_active: a stale or + # unsolicited confirm must not extend the transaction TTL nor reset idle budgets + if ref.obj_confirmed(requester, status, nonce): + ref.mark_active() + finally: + ref.tx.end_op() + return make_reply(ReturnCode.OK) @classmethod def _monitor_tx(cls): while True: now = time.time() + + # Per-receiver budget enforcement runs OUTSIDE _tx_lock: finalizing a + # budget-failed receiver fires user callbacks (downloaded_to_one/all), which must + # never run under the global lock. A budget failure recorded here flips + # is_finished() so the classification pass below resolves the tx immediately. + with cls._tx_lock: + budget_txs = [tx for tx in cls._tx_table.values() if tx.has_receiver_budgets] + for tx in budget_txs: + with cls._tx_lock: + # dead or settling tx: skip (begin_op makes the table check binding) + live = cls._tx_table.get(tx.tid) is tx and tx.begin_op() + if not live: + continue + try: + tx.enforce_receiver_budgets(now) + except Exception as ex: + cls._logger.error( + f"error enforcing receiver budgets for tx {tx.tid}: {secure_format_exception(ex)}" + ) + finally: + tx.end_op() + expired_tx = [] finished_tx = [] with cls._tx_lock: @@ -949,16 +1599,19 @@ def _monitor_tx(cls): cls._expire_finished_refs(now) cls._expire_outcomes(now) + cls._reap_termination_markers() for tx in expired_tx: tx.transaction_done( TransactionDoneStatus.TIMEOUT, on_outcome=functools.partial(cls._record_outcome, tx=tx) ) + cls._sync_termination_marker(tx) for tx in finished_tx: tx.transaction_done( TransactionDoneStatus.FINISHED, on_outcome=functools.partial(cls._record_outcome, tx=tx) ) + cls._sync_termination_marker(tx) time.sleep(5.0) @@ -1055,6 +1708,39 @@ def download_object( # On retry, resend the same state so producer re-generates the same chunk. current_state = None + # Receiver-confirmed completion: we advertise the capability on every request (when the + # kill-switch is on) and learn from each reply whether the producer consumes confirmations. + confirm_enabled = _receiver_confirm_enabled() + producer_expects_confirm = False + confirm_nonce = None + + def _send_confirm(receiver_truth: str): + # wire contract: a confirmation is sent ONLY after a producer-served terminal reply + # (EOF/ERROR) -- the producer accepts a confirm only against its pending provisional + # serve, so mid-stream failure exits do not confirm (budgets/TTL handle those) + if not (confirm_enabled and producer_expects_confirm): + return + try: + # fire-and-forget by design: a lost confirmation is backstopped producer-side by + # per-receiver budgets / the transaction timeout, failing closed + cell.fire_and_forget( + channel=OBJ_DOWNLOADER_CHANNEL, + topic=OBJ_DOWNLOADER_TOPIC, + targets=from_fqcn, + message=new_cell_message( + headers={}, + payload={ + _PropKey.REF_ID: ref_id, + _PropKey.CONFIRM: receiver_truth, + _PropKey.CONFIRM_NONCE: confirm_nonce, + }, + ), + secure=secure, + optional=optional, + ) + except Exception as ex: + logger.warning(f"failed to send download confirmation for ref={ref_id}: {secure_format_exception(ex)}") + def _emit_progress(state: str, force: bool = False): nonlocal progress_sequence, last_progress_emit_time if not progress_cb: @@ -1084,6 +1770,8 @@ def _emit_progress(state: str, force: bool = False): # Build a fresh request each iteration (including retries) # to avoid re-encoding an already-encoded message. request_payload = {_PropKey.REF_ID: ref_id} + if confirm_enabled: + request_payload[_PropKey.CONFIRM_CAPABLE] = True if current_state is not None: request_payload[_PropKey.STATE] = current_state request = new_cell_message(headers={}, payload=request_payload) @@ -1151,6 +1839,9 @@ def _emit_progress(state: str, force: bool = False): payload = reply.payload assert isinstance(payload, dict) + if payload.get(_PropKey.CONFIRM_EXPECTED): + producer_expects_confirm = True + confirm_nonce = payload.get(_PropKey.CONFIRM_NONCE) status = payload.get(_PropKey.STATUS) if status == ProduceRC.EOF: elapsed = time.time() - download_start @@ -1159,10 +1850,20 @@ def _emit_progress(state: str, force: bool = False): f"[client] download ref={ref_id} done: elapsed={elapsed:.2f}s " f"size={size_mb:.1f}MB ({total_bytes:,} bytes)" ) - consumer.download_completed(ref_id) + try: + consumer.download_completed(ref_id) + except Exception: + # receiver-side finalization failed AFTER the last chunk (e.g. disk-offload + # finalize): exactly what receiver-confirmed completion exists to surface -- + # the producer must not certify this receiver on its served EOF + _send_confirm(DownloadStatus.FAILED) + _emit_progress("failed", force=True) + raise + _send_confirm(DownloadStatus.SUCCESS) _emit_progress("completed", force=True) return elif status == ProduceRC.ERROR: + _send_confirm(DownloadStatus.FAILED) consumer.download_failed(ref_id, f"producer error after {duration} secs") _emit_progress("failed", force=True) return diff --git a/nvflare/fuel/f3/streaming/obj_downloader.py b/nvflare/fuel/f3/streaming/obj_downloader.py index 7fdb7ad8ca..c2b1ea68f5 100644 --- a/nvflare/fuel/f3/streaming/obj_downloader.py +++ b/nvflare/fuel/f3/streaming/obj_downloader.py @@ -28,6 +28,10 @@ def __init__( progress_cb=None, progress_interval: float = 30.0, outcome_cb=None, + receiver_ids=None, + min_receivers=None, + receiver_acquire_timeout=None, + receiver_idle_timeout=None, **cb_kwargs, ): """Constructor of ObjectDownloader. @@ -67,6 +71,10 @@ def __init__( progress_cb=progress_cb, progress_interval=progress_interval, outcome_cb=outcome_cb, + receiver_ids=receiver_ids, + min_receivers=min_receivers, + receiver_acquire_timeout=receiver_acquire_timeout, + receiver_idle_timeout=receiver_idle_timeout, **cb_kwargs, ) @@ -87,6 +95,15 @@ def add_object(self, obj: Downloadable, ref_id=None) -> str: ) return rid + def get_waiter(self): + """Returns the awaitable facade (TransferWaiter) over this transaction's terminal outcome. + + waiter.wait(timeout) blocks until the aggregate TransferOutcome is recorded -- + COMPLETED only when every expected receiver succeeded. This is the primitive + upper layers use to guarantee "send returns only when the payload is delivered". + """ + return DownloadService.get_transfer_waiter(self.tx_id) + def delete_transaction(self): """Delete the download transaction forcefully. You call this method only if you want to stop the downloading process prematurely. diff --git a/nvflare/fuel/f3/streaming/transfer_outcome.py b/nvflare/fuel/f3/streaming/transfer_outcome.py index 45ed3736ba..f67d0ab61a 100644 --- a/nvflare/fuel/f3/streaming/transfer_outcome.py +++ b/nvflare/fuel/f3/streaming/transfer_outcome.py @@ -31,20 +31,24 @@ Outcome status values reuse the TransferProgressState terminal vocabulary (completed / failed / aborted) rather than introducing another status set. -Known limits, resolved by later PRs of the same design (see -docs/design/client_api_execution_modes_plan.md): -- per-receiver statuses are producer-served (recorded when produce() returns EOF), - not receiver-confirmed — a receiver-side finalization failure after the last chunk - is not visible here until receiver-confirmed completion (plan PR F3-2) lands; -- num_receivers is a count without receiver identity — expected-receiver identity - checks arrive with per-receiver budgets (plan PR F3-3); +Semantics with the full payload layer (receiver-confirmed completion, per-receiver +budgets, awaitable transfer facade): +- per-receiver statuses are receiver-confirmed where the receiver supports it (a served + EOF is provisional until the receiver confirms its finalization succeeded); legacy + receivers remain producer-served — both skews and the runtime kill-switch degrade to + producer-served semantics (download_service.py, receiver-confirmed completion); +- expected receiver identities and per-receiver acquire/idle budgets bound the outcome's + resolution time (a stalled or never-pulling receiver is finalized FAILED without + waiting the whole-transaction TTL); min_receivers surfaces the optional k-of-N quorum + via quorum_met while `completed` stays the strict all-receivers certificate; - the outcome covers the refs present at termination; adding objects to a transaction after receivers already finished the earlier ones is not supported. """ import time from dataclasses import dataclass -from typing import Dict, List, Optional +from types import MappingProxyType +from typing import Mapping, Optional, Sequence, Tuple from nvflare.fuel.f3.streaming.transfer_progress import TransferProgressState @@ -85,19 +89,31 @@ class TransferOutcomeReason: DELETED = "deleted" UNKNOWN_RECEIVER_COUNT = "unknown_receiver_count" UNKNOWN_DONE_STATUS = "unknown_done_status" + # the verdict computation itself failed: settlement recorded a fail-closed + # placeholder so ownership is consumed and waiters resolve + COMPUTATION_FAILED = "computation_failed" @dataclass(frozen=True) class RefOutcome: """Per-object terminal outcome. - receiver_statuses maps receiver FQCN to a DownloadStatus value (success / failed) - as recorded by the producer side. Treat the dict as read-only: the same instance - is shared with every consumer of the outcome. + receiver_statuses maps receiver FQCN to a DownloadStatus value (success / failed) -- + receiver-confirmed where the receiver supports it, producer-served for legacy peers + or when the receiver-confirm kill-switch is off, and budget-FAILED for receivers that + exhausted their acquire/idle budget. It is deep-frozen at construction (a + MappingProxyType over a private copy): the same instance is recorded in the + service outcome table and handed to outcome_cb consumers, so a callback must + not be able to mutate the recorded per-receiver truth. If outcomes ever cross + a process boundary, the serializer must materialize it (dict(...)). """ ref_id: str - receiver_statuses: Dict[str, str] + receiver_statuses: Mapping[str, str] + + def __post_init__(self): + # frozen=True only blocks attribute rebinding; freeze the container too + object.__setattr__(self, "receiver_statuses", MappingProxyType(dict(self.receiver_statuses))) @dataclass(frozen=True) @@ -116,20 +132,68 @@ class TransferOutcome: reason: str # a TransferOutcomeReason value done_status: str # the raw TransactionDoneStatus value num_receivers: int - refs: List[RefOutcome] + refs: Tuple[RefOutcome, ...] timestamp: float + # optional k-of-N quorum declared by the workflow: informational for quorum_met; + # `completed` deliberately stays the strict all-receivers certificate + min_receivers: Optional[int] = None + # expected receiver identities declared by the workflow. When present, completion and + # quorum are judged against THESE identities: a status from an unexpected receiver can + # neither complete the transfer nor count toward the quorum. + receiver_ids: Optional[Tuple[str, ...]] = None + + def __post_init__(self): + # frozen=True only blocks attribute rebinding; freeze the containers too so + # outcome_cb consumers cannot mutate the recorded outcome + object.__setattr__(self, "refs", tuple(self.refs)) + if self.receiver_ids is not None: + object.__setattr__(self, "receiver_ids", tuple(self.receiver_ids)) @property def completed(self) -> bool: return self.status == TransferProgressState.COMPLETED + @property + def quorum_met(self) -> bool: + """True if every ref reached at least min_receivers confirmed successes. + + The k-of-N surface for fan-out workflows (min_responses-style policies): `completed` + stays the strict all-receivers certificate; a workflow that accepts partial fan-out + checks quorum_met (or thresholds refs itself). Falls back to `completed` when no + min_receivers was declared; fails closed with no refs. + """ + if self.min_receivers is None: + return self.completed + if not self.refs: + return False + # a receiver counts toward the quorum only if it succeeded on EVERY ref: counting + # per-ref successes independently would let different receiver subsets satisfy each + # ref while no single receiver holds the complete payload + quorum_receivers = None + for r in self.refs: + ref_successes = {rcv for rcv, v in r.receiver_statuses.items() if v == DownloadStatus.SUCCESS} + quorum_receivers = ref_successes if quorum_receivers is None else quorum_receivers & ref_successes + if self.receiver_ids is not None: + # only declared receivers count toward the quorum + quorum_receivers &= set(self.receiver_ids) + return len(quorum_receivers) >= self.min_receivers + def expired(self, now: float, ttl: float) -> bool: return now - self.timestamp > ttl -def _all_receivers_succeeded(num_receivers: int, refs: List[RefOutcome]) -> bool: +def _all_receivers_succeeded(num_receivers: int, refs: Sequence[RefOutcome], receiver_ids=None) -> bool: if num_receivers <= 0 or not refs: return False + if receiver_ids: + # identity mode: every DECLARED receiver must have succeeded on every ref. + # Statuses from unexpected receivers are ignored -- they can never certify + # a transfer that a declared receiver did not actually get. + for r in refs: + for expected in receiver_ids: + if r.receiver_statuses.get(expected) != DownloadStatus.SUCCESS: + return False + return True for r in refs: if len(r.receiver_statuses) < num_receivers: return False @@ -149,8 +213,10 @@ def compute_transfer_outcome( tx_id: str, done_status: str, num_receivers: int, - refs: List[RefOutcome], + refs: Sequence[RefOutcome], timestamp: Optional[float] = None, + min_receivers: Optional[int] = None, + receiver_ids=None, ) -> TransferOutcome: """Compute the aggregate terminal outcome for a terminated transaction. @@ -175,7 +241,7 @@ def compute_transfer_outcome( if done_status not in _KNOWN_DONE_STATUSES: status, reason = TransferProgressState.FAILED, TransferOutcomeReason.UNKNOWN_DONE_STATUS - elif _all_receivers_succeeded(num_receivers, refs): + elif _all_receivers_succeeded(num_receivers, refs, receiver_ids): status, reason = TransferProgressState.COMPLETED, TransferOutcomeReason.ALL_RECEIVERS_SUCCEEDED elif done_status == TransactionDoneStatus.DELETED: status, reason = TransferProgressState.ABORTED, TransferOutcomeReason.DELETED @@ -200,4 +266,6 @@ def compute_transfer_outcome( num_receivers=num_receivers, refs=refs, timestamp=timestamp, + min_receivers=min_receivers, + receiver_ids=receiver_ids, ) diff --git a/nvflare/fuel/utils/fobs/decomposers/via_downloader.py b/nvflare/fuel/utils/fobs/decomposers/via_downloader.py index a90e1caddb..b461e3319d 100644 --- a/nvflare/fuel/utils/fobs/decomposers/via_downloader.py +++ b/nvflare/fuel/utils/fobs/decomposers/via_downloader.py @@ -380,7 +380,9 @@ def decompose(self, target: Any, manager: DatumManager = None) -> Any: self.logger.debug(f"ViaDownloader: created ref for target {target_id}: {item_id}") return {EncKey.TYPE: EncType.REF, EncKey.DATA: item_id} - def _create_downloader(self, fobs_ctx: dict, progress_cb=None, timeout_override=None, num_receivers_override=None): + def _create_downloader( + self, fobs_ctx: dict, progress_cb=None, timeout_override=None, num_receivers_override=None, receiver_ids=None + ): # Transaction lifecycle is managed solely by _monitor_tx() (download_service.py). # We deliberately do NOT subscribe to msg_root deletion here. The msg_root is # deleted as soon as all blobs are delivered, but blob_cb fires asynchronously — @@ -445,6 +447,9 @@ def _create_downloader(self, fobs_ctx: dict, progress_cb=None, timeout_override= transaction_done_cb=on_complete_cb, progress_cb=progress_cb, progress_interval=RESULT_UPLOAD_PROGRESS_INTERVAL, + # expected receiver identities: enables the transaction's per-receiver + # acquire budget; None when any identity is unknown + receiver_ids=receiver_ids, ) return downloader @@ -603,11 +608,15 @@ def _finalize_download_tx(self, mgr: DatumManager): progress_context = fobs_ctx.get(RESULT_UPLOAD_PROGRESS_CTX_KEY) or {} timeout_override = progress_context.get(ResultUploadProgressContextKey.STREAMING_IDLE_TIMEOUT) + # forward receiver identities to the transaction only when every identity is + # actually known (a (None,) placeholder means unknown-single-receiver) + known_receiver_ids = receiver_ids if receiver_ids and all(r is not None for r in receiver_ids) else None downloader = self._create_downloader( fobs_ctx, progress_cb=progress_cb if progress_trackable else None, timeout_override=timeout_override, num_receivers_override=download_num_receivers, + receiver_ids=known_receiver_ids, ) if downloader is None: self.logger.warning("download transaction was not created because FOBS context has no cell") diff --git a/tests/unit_test/client/cell/defs_test.py b/tests/unit_test/client/cell/defs_test.py index b5644a09f2..0b16c0244d 100644 --- a/tests/unit_test/client/cell/defs_test.py +++ b/tests/unit_test/client/cell/defs_test.py @@ -33,6 +33,7 @@ def test_expected_topics_present(self): "HELLO_REJECTED", "TASK_READY", "TASK_ACCEPTED", + "TASK_PAYLOAD_READY", "TASK_FAILED", "RESULT_READY", "RESULT_ACCEPTED", @@ -77,6 +78,7 @@ def test_topic_wire_values(self): "HELLO_REJECTED": "client_api.hello_rejected", "TASK_READY": "client_api.task_ready", "TASK_ACCEPTED": "client_api.task_accepted", + "TASK_PAYLOAD_READY": "client_api.task_payload_ready", "TASK_FAILED": "client_api.task_failed", "RESULT_READY": "client_api.result_ready", "RESULT_ACCEPTED": "client_api.result_accepted", diff --git a/tests/unit_test/fuel/f3/streaming/conftest.py b/tests/unit_test/fuel/f3/streaming/conftest.py new file mode 100644 index 0000000000..fa33b66f83 --- /dev/null +++ b/tests/unit_test/fuel/f3/streaming/conftest.py @@ -0,0 +1,35 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. 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. +"""DIRECTORY-WIDE fixtures: the autouse fixture below applies to EVERY test in +tests/unit_test/fuel/f3/streaming/, including tests that never touch confirms. +A new test here never exercises the real ConfigService read of the confirm +kill-switch; patch _receiver_confirm_cached explicitly if you need the legacy +(confirm-off) path or the real config read.""" + +from unittest.mock import patch + +import pytest + +from nvflare.fuel.f3.streaming import download_service as ds_module + + +@pytest.fixture(autouse=True) +def confirm_switch_on(): + """Pin the receiver-confirm kill-switch ON for all streaming tests. + + Individual tests patch it OFF explicitly; legacy-path tests are unaffected (their + requests carry no capability key, so the switch is never consulted for them). + """ + with patch.object(ds_module, "_receiver_confirm_cached", True): + yield diff --git a/tests/unit_test/fuel/f3/streaming/download_service_test.py b/tests/unit_test/fuel/f3/streaming/download_service_test.py index 67c14d88ac..8cdff791d8 100644 --- a/tests/unit_test/fuel/f3/streaming/download_service_test.py +++ b/tests/unit_test/fuel/f3/streaming/download_service_test.py @@ -13,6 +13,7 @@ # limitations under the License. import threading +import time from typing import Any from unittest.mock import Mock, patch @@ -28,6 +29,7 @@ ProduceRC, TransactionDoneStatus, ) +from nvflare.fuel.f3.streaming.transfer_outcome import compute_transfer_outcome from nvflare.fuel.utils.network_utils import get_open_ports from tests.unit_test.fuel.f3.streaming.download_test_utils import ( MockDownloadable, @@ -476,50 +478,40 @@ def register_request_cb(self, **kwargs): assert list(service._initialized_cells.keys()) == [] - def test_initialize_reenable_holds_outcome_lock(self): - """The _accept_outcomes re-enable in _initialize() must happen under _outcome_lock. + def test_new_transaction_takes_ownership_before_monitor_visible(self): + """A tx must own its outcome slot before it appears in _tx_table. - This guards lock-discipline uniformity only: the flag is read under _outcome_lock in - _record_outcome(), so its write must be too (it previously ran under _init_lock only). - This is NOT what closes the stale-outcome race -- that is the live-incarnation guard - in _record_outcome(), covered by test_stale_outcome_dropped_after_incarnation_cleared. + The monitor discovers transactions through _tx_table. If a tx were inserted + there first, a monitor tick landing in the window before ownership is taken + could terminate it: its terminal outcome would be dropped by the owner guard + (a terminated-but-unknown gap), and new_transaction would then register a + dead owner entry that nothing ever pops. """ service = _make_isolated_download_service() service._tx_monitor = object() # avoid starting a real monitor thread - class FakeCell: - def register_request_cb(self, **kwargs): - pass + registered_at_insert = [] - cell = FakeCell() - service._accept_outcomes = False # as left by a prior shutdown() - - init_done = threading.Event() - - def run_init(): - service._initialize(cell) - init_done.set() - - with service._outcome_lock: - t = threading.Thread(target=run_init) - t.start() - # while _outcome_lock is held, the guarded re-enable must not complete - assert not init_done.wait(0.2) - assert service._accept_outcomes is False - - t.join(2.0) - assert init_done.is_set() - assert service._accept_outcomes is True - - def test_stale_outcome_dropped_after_incarnation_cleared(self): - """A terminal outcome for a transaction whose incarnation is no longer live must drop. - - This is the invariant that actually closes the cross-lifecycle stale-outcome race: - _record_outcome() records only for the live registered incarnation (current is tx). - After shutdown() clears _tx_incarnations, a callback that blocked on _outcome_lock - during shutdown can win the lock afterward and observe _accept_outcomes re-enabled by - a subsequent _initialize() -- so the _accept_outcomes flag alone does not stop it. The - live-incarnation guard does: with no live incarnation for the tid, the outcome drops. + class MonitorVisibilityDict(dict): + def __setitem__(self, key, value): + registered_at_insert.append(key in service._outcome_owners) + super().__setitem__(key, value) + + service._tx_table = MonitorVisibilityDict() + + tx_id = service.new_transaction(cell=Mock(), timeout=10.0, num_receivers=1) + + assert registered_at_insert == [True] + assert service._outcome_owners[tx_id] is service._tx_table[tx_id] + + def test_stale_outcome_dropped_after_ownership_cleared(self): + """A terminal outcome for a transaction that no longer owns its tx_id must drop. + + This is the invariant that closes the cross-lifecycle stale-outcome race: + _record_outcome() records only for the transaction that owns the outcome slot. + A recorder that blocked on _outcome_lock while shutdown() cleared the tables can + win the lock after a subsequent _initialize(); no longer owning the slot, its + pre-shutdown outcome drops instead of repopulating the cleared table. """ service = _make_isolated_download_service() @@ -529,23 +521,57 @@ def test_stale_outcome_dropped_after_incarnation_cleared(self): old_outcome.tx_id = "tx-old" old_outcome.expired.return_value = False - # shutdown() cleared incarnations; a later _initialize() re-enabled recording - service._tx_incarnations.clear() - service._accept_outcomes = True + # shutdown() cleared ownership; recording is gated by owner identity alone + service._outcome_owners.clear() # the late callback for the old, now-unregistered transaction must be dropped service._record_outcome(old_outcome, tx=old_tx) assert service.get_transaction_outcome("tx-old") is None - # sanity: a live registered incarnation still records (guard does not over-drop) + # sanity: the owning transaction still records (guard does not over-drop) live_tx = Mock() live_tx.tid = "tx-live" - live_outcome = Mock() - live_outcome.tx_id = "tx-live" - live_outcome.expired.return_value = False - service._tx_incarnations["tx-live"] = live_tx + live_outcome = compute_transfer_outcome("tx-live", TransactionDoneStatus.FINISHED, 1, [], time.time()) + service._outcome_owners["tx-live"] = live_tx service._record_outcome(live_outcome, tx=live_tx) - assert service.get_transaction_outcome("tx-live") is live_outcome + recorded = service.get_transaction_outcome("tx-live") + # recording re-stamps the receipt (TTL starts at recording), so compare identity-free + assert recorded is not None and recorded.tx_id == "tx-live" and recorded.done_status == live_outcome.done_status + + def test_raising_download_callbacks_do_not_break_serving_path(self): + """Raising downloaded_to_one/downloaded_to_all must not propagate into serving. + + A raising downloaded_to_one on the chunk-serving path would lose the EOF reply + for that attempt, and -- because the _downloaded_to_all_called latch is set + before the callbacks run and is never retried -- would permanently skip + downloaded_to_all. Both callbacks are guarded like the terminal callbacks in + transaction_done: the exception is logged, serving and the all-receivers-done + notification proceed. + """ + service = _make_isolated_download_service() + service._tx_monitor = object() # avoid starting a real monitor thread + + calls = [] + + class RaisingDownloadable(MockDownloadable): + def downloaded_to_one(self, to_receiver: str, status: str): + calls.append(("one", to_receiver, status)) + raise RuntimeError("user callback failure") + + def downloaded_to_all(self): + calls.append(("all",)) + raise RuntimeError("user callback failure") + + tx_id = service.new_transaction(cell=Mock(), timeout=10.0, num_receivers=1) + obj = RaisingDownloadable([b"chunk"]) + rid = service.add_object(tx_id, obj) + ref = service._ref_table[rid] + + # must not raise; downloaded_to_all still fires after downloaded_to_one raised + ref.obj_downloaded("r1", DownloadStatus.SUCCESS) + + assert calls == [("one", "r1", DownloadStatus.SUCCESS), ("all",)] + assert ref.snapshot_receiver_statuses() == {"r1": DownloadStatus.SUCCESS} def test_get_transaction_id_from_ref_id(self, cell): """Test retrieving transaction ID from reference ID.""" diff --git a/tests/unit_test/fuel/f3/streaming/download_test_utils.py b/tests/unit_test/fuel/f3/streaming/download_test_utils.py index 1e218027da..e5c806934e 100644 --- a/tests/unit_test/fuel/f3/streaming/download_test_utils.py +++ b/tests/unit_test/fuel/f3/streaming/download_test_utils.py @@ -26,7 +26,9 @@ import pytest -from nvflare.fuel.f3.streaming.download_service import Downloadable, DownloadService, ProduceRC +from nvflare.fuel.f3.cellnet.defs import MessageHeaderKey +from nvflare.fuel.f3.cellnet.utils import new_cell_message +from nvflare.fuel.f3.streaming.download_service import Downloadable, DownloadService, ProduceRC, _PropKey class MockDownloadable(Downloadable): @@ -83,9 +85,10 @@ class IsolatedDownloadService(DownloadService): _ref_table = {} _finished_refs = {} _tx_outcomes = {} - _tx_incarnations = {} + _outcome_owners = {} + _tx_waiters = {} + _terminating_txs = {} _outcome_lock = threading.Lock() - _accept_outcomes = True _logger = Mock() _tx_lock = threading.Lock() _initialized_cells = weakref.WeakKeyDictionary() @@ -93,6 +96,52 @@ class IsolatedDownloadService(DownloadService): return IsolatedDownloadService +def make_service_no_monitor(): + """An isolated service with the real monitor thread suppressed. + + Not confirm-specific: use for any test that must not race the real 5s monitor + loop (budgets, confirms, waiters, lifecycle); drive monitor passes explicitly + via run_monitor_once.""" + service = make_isolated_download_service() + service._tx_monitor = Mock() + return service + + +def pull_request(rid, requester, confirm_capable=False, state=None): + """Builds one downloader pull request as it appears on the wire.""" + payload = {_PropKey.REF_ID: rid} + if confirm_capable: + payload[_PropKey.CONFIRM_CAPABLE] = True + if state is not None: + payload[_PropKey.STATE] = state + return new_cell_message(headers={MessageHeaderKey.ORIGIN: requester}, payload=payload) + + +def confirm_request(rid, requester, status, nonce=None): + """Builds one receiver-confirmation message as it appears on the wire.""" + payload = {_PropKey.REF_ID: rid, _PropKey.CONFIRM: status} + if nonce is not None: + payload[_PropKey.CONFIRM_NONCE] = nonce + return new_cell_message(headers={MessageHeaderKey.ORIGIN: requester}, payload=payload) + + +def serve_nonce(terminal_reply): + """The per-serve nonce a confirmation must echo (from the producer's terminal reply).""" + return terminal_reply.payload.get(_PropKey.CONFIRM_NONCE) + + +def pull_to_terminal(service, rid, requester, confirm_capable=False): + """Drives the pull loop for one receiver until the producer serves a terminal status.""" + state = None + for _ in range(50): + reply = service._handle_download(pull_request(rid, requester, confirm_capable=confirm_capable, state=state)) + status = reply.payload.get(_PropKey.STATUS) + if status in (ProduceRC.EOF, ProduceRC.ERROR): + return reply + state = reply.payload.get(_PropKey.STATE) + raise AssertionError("pull loop never reached a terminal status") + + def run_monitor_once(service_cls, now): from nvflare.fuel.f3.streaming import download_service as download_service_module diff --git a/tests/unit_test/fuel/f3/streaming/receiver_budget_test.py b/tests/unit_test/fuel/f3/streaming/receiver_budget_test.py new file mode 100644 index 0000000000..5dc0ebab0d --- /dev/null +++ b/tests/unit_test/fuel/f3/streaming/receiver_budget_test.py @@ -0,0 +1,472 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. 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. + +"""Tests for per-(transfer, receiver) acquire/idle budgets and the quorum surface. + +A receiver that exhausts its acquire budget (never issued its first pull) or idle budget +(stopped making requests) is finalized FAILED for the aggregate outcome on a monitor pass -- +without waiting for the whole-transaction TTL, and without a live receiver masking a stalled +one behind the tx-wide activity timestamp. Budgets are per-transaction opt-in (config-var +defaults); disabled budgets preserve today's TTL-only behavior exactly. +""" + +import time +from unittest.mock import Mock, patch + +import pytest + +from nvflare.fuel.f3.streaming import download_service as ds_module +from nvflare.fuel.f3.streaming.download_service import DownloadStatus +from nvflare.fuel.f3.streaming.transfer_outcome import TransferOutcomeReason, compute_transfer_outcome +from tests.unit_test.fuel.f3.streaming.download_test_utils import MockDownloadable, confirm_request +from tests.unit_test.fuel.f3.streaming.download_test_utils import make_service_no_monitor as _make_service +from tests.unit_test.fuel.f3.streaming.download_test_utils import pull_request +from tests.unit_test.fuel.f3.streaming.download_test_utils import pull_to_terminal as _pull_to_terminal +from tests.unit_test.fuel.f3.streaming.download_test_utils import run_monitor_once, serve_nonce + + +def _new_tx(service, chunks=1, **tx_kwargs): + tx_kwargs.setdefault("timeout", 1000.0) # the whole-tx TTL is deliberately huge: + tx_kwargs.setdefault("num_receivers", 1) # budgets, not the TTL, must resolve these tests + tx_id = service.new_transaction(cell=Mock(), **tx_kwargs) + obj = MockDownloadable([b"chunk"] * chunks) + rid = service.add_object(tx_id, obj) + return tx_id, rid + + +class TestIdleBudget: + def test_stalled_receiver_fails_without_waiting_tx_ttl(self): + service = _make_service() + tx_id, rid = _new_tx(service, chunks=3, num_receivers=2, receiver_idle_timeout=5.0) + + _pull_to_terminal(service, rid, "healthy") # legacy receiver, final at serve + service._handle_download(pull_request(rid, "stalled")) # one pull, then silence + + # a monitor pass past the idle budget (but far inside the 1000s TTL) + run_monitor_once(service, now=time.time() + 30.0) + + outcome = service.get_transaction_outcome(tx_id) + assert outcome is not None, "budget failure must resolve the tx on this pass, not at TTL" + assert not outcome.completed + assert outcome.reason == TransferOutcomeReason.RECEIVER_FAILED + statuses = outcome.refs[0].receiver_statuses + assert statuses["healthy"] == DownloadStatus.SUCCESS + assert statuses["stalled"] == DownloadStatus.FAILED + + def test_lost_confirmation_is_bounded_by_idle_budget(self): + # fire-and-forget confirms can be lost: the receiver stops making requests after its + # served EOF, so its idle budget finalizes it FAILED in bounded time (fail-closed) + service = _make_service() + tx_id, rid = _new_tx(service, receiver_idle_timeout=5.0) + _pull_to_terminal(service, rid, "r1", confirm_capable=True) # provisional, confirm never arrives + + run_monitor_once(service, now=time.time() + 30.0) + + outcome = service.get_transaction_outcome(tx_id) + assert outcome is not None + assert not outcome.completed + assert outcome.refs[0].receiver_statuses == {"r1": DownloadStatus.FAILED} + # the provisional record was cleaned up, not leaked + assert service._tx_outcomes # outcome recorded; ref is gone with the finished tx + + def test_active_receiver_is_not_idle_failed(self): + service = _make_service() + tx_id, rid = _new_tx(service, chunks=3, receiver_idle_timeout=5.0) + reply = service._handle_download(pull_request(rid, "r1")) + ref = service._ref_table[rid] + tx = service._tx_table[tx_id] + # receiver keeps making requests: refresh activity to "now" as seen by the monitor + future = time.time() + 30.0 + with tx._stats_lock: + tx._receiver_last_active["r1"] = future - 1.0 + + run_monitor_once(service, now=future) + + assert service.get_transaction_outcome(tx_id) is None # still live + assert ref.snapshot_receiver_statuses() == {} + + +class TestAcquireBudget: + def test_never_pulling_receiver_fails_at_acquire_deadline(self): + service = _make_service() + tx_id, rid = _new_tx( + service, + num_receivers=0, # derived from receiver_ids + receiver_ids=("r1", "r2"), + receiver_acquire_timeout=5.0, + min_receivers=1, + ) + _pull_to_terminal(service, rid, "r1") # r2 never shows up + + run_monitor_once(service, now=time.time() + 30.0) + + outcome = service.get_transaction_outcome(tx_id) + assert outcome is not None + statuses = outcome.refs[0].receiver_statuses + assert statuses == {"r1": DownloadStatus.SUCCESS, "r2": DownloadStatus.FAILED} + # strict certificate fails; the declared k-of-N quorum is satisfied + assert not outcome.completed + assert outcome.quorum_met + assert outcome.min_receivers == 1 + + def test_receiver_with_activity_is_not_acquire_failed(self): + service = _make_service() + tx_id, rid = _new_tx(service, chunks=3, num_receivers=0, receiver_ids=("r1",), receiver_acquire_timeout=5.0) + service._handle_download(pull_request(rid, "r1")) # first pull happened: acquire satisfied + + run_monitor_once(service, now=time.time() + 30.0) + + # no idle budget configured, so the slow-but-acquired receiver is left alone + assert service.get_transaction_outcome(tx_id) is None + assert service._ref_table[rid].snapshot_receiver_statuses() == {} + + def test_acquire_budget_needs_receiver_identities(self): + # without receiver_ids there is no one to hold to the acquire deadline + service = _make_service() + tx_id, rid = _new_tx(service, num_receivers=2, receiver_acquire_timeout=5.0) + + run_monitor_once(service, now=time.time() + 30.0) + + assert service.get_transaction_outcome(tx_id) is None + assert service._ref_table[rid].snapshot_receiver_statuses() == {} + + +class TestBudgetSemantics: + def test_budgets_disabled_preserve_ttl_only_behavior(self): + service = _make_service() + tx_id, rid = _new_tx(service, num_receivers=2) + service._handle_download(pull_request(rid, "r1")) + + run_monitor_once(service, now=time.time() + 500.0) # inside the 1000s TTL + + assert service.get_transaction_outcome(tx_id) is None + assert service._ref_table[rid].snapshot_receiver_statuses() == {} + + def test_activity_tracked_without_progress_cb(self): + service = _make_service() + tx_id, rid = _new_tx(service, chunks=2) + service._handle_download(pull_request(rid, "r1")) + + activity = service._ref_table[rid].snapshot_receiver_activity() + assert "r1" in activity + + def test_budget_failed_receiver_is_final_even_against_late_confirm(self): + service = _make_service() + tx_id, rid = _new_tx(service, num_receivers=2, receiver_idle_timeout=5.0) + terminal = _pull_to_terminal(service, rid, "r1", confirm_capable=True) + ref = service._ref_table[rid] + + run_monitor_once(service, now=time.time() + 30.0) # idle budget fires + assert ref.snapshot_receiver_statuses()["r1"] == DownloadStatus.FAILED + assert ref.snapshot_pending_confirms() == {} + + # the straggler confirmation (correct nonce and all) cannot resurrect it + assert ref.obj_confirmed("r1", DownloadStatus.SUCCESS, serve_nonce(terminal)) is False + assert ref.snapshot_receiver_statuses()["r1"] == DownloadStatus.FAILED + + +class TestMultiRefAcquisition: + def test_sequential_multi_ref_receiver_is_not_acquire_failed(self): + # transaction-level acquisition: a healthy receiver pulling ref 1 must not be + # acquire-failed on ref 2 it has not reached yet + service = _make_service() + tx_id = service.new_transaction( + cell=Mock(), timeout=1000.0, num_receivers=0, receiver_ids=("r1",), receiver_acquire_timeout=5.0 + ) + rid1 = service.add_object(tx_id, MockDownloadable([b"chunk"] * 3)) + rid2 = service.add_object(tx_id, MockDownloadable([b"chunk"] * 3)) + service._handle_download(pull_request(rid1, "r1")) # busy on ref 1, has not touched ref 2 + + run_monitor_once(service, now=time.time() + 30.0) + + assert service.get_transaction_outcome(tx_id) is None # still live + assert service._ref_table[rid2].snapshot_receiver_statuses() == {} + + def test_confirmed_receiver_wins_over_stale_budget_snapshot(self): + # truth-wins re-check: a receiver finalized (confirmed) between the budget snapshot + # and enforcement must not be flipped to FAILED, and no failure is reported for it + service = _make_service() + tx_id, rid = _new_tx(service, receiver_idle_timeout=5.0) + terminal = _pull_to_terminal(service, rid, "r1", confirm_capable=True) + ref = service._ref_table[rid] + service._handle_download(confirm_request(rid, "r1", DownloadStatus.SUCCESS, serve_nonce(terminal))) + + tx = service._tx_table[tx_id] + with tx._stats_lock: + tx_last_active = dict(tx._receiver_last_active) + enforced = ref.enforce_budgets( + time.time() + 30.0, None, 5.0, None, tx_acquired={"r1"}, tx_last_active=tx_last_active + ) + + assert enforced == [] + assert ref.snapshot_receiver_statuses() == {"r1": DownloadStatus.SUCCESS} + + +class TestMultiRefIdleEscape: + def test_receiver_idle_after_finishing_one_ref_fails_sibling_ref_in_bounded_time(self): + # P2 pin (reproduced by review): after finishing ref 1, the receiver is tx-acquired + # (exempt from ref 2's acquire budget) and has no per-ref timestamp on ref 2 -- with + # per-ref idle it escaped BOTH budgets and pinned the producer to the full TTL. + # Idle is judged on transaction-level activity, so its silence now fails ref 2. + service = _make_service() + tx_id = service.new_transaction( + cell=Mock(), timeout=1000.0, num_receivers=0, receiver_ids=("r1",), receiver_idle_timeout=5.0 + ) + rid1 = service.add_object(tx_id, MockDownloadable([b"chunk"])) + rid2 = service.add_object(tx_id, MockDownloadable([b"chunk"])) + _pull_to_terminal(service, rid1, "r1") # finishes ref 1 (legacy final), never touches ref 2 + + run_monitor_once(service, now=time.time() + 30.0) + + outcome = service.get_transaction_outcome(tx_id) + assert outcome is not None, "must resolve via the idle budget, not the 1000s TTL" + assert not outcome.completed + statuses = {r.ref_id: r.receiver_statuses for r in outcome.refs} + assert statuses[rid1] == {"r1": DownloadStatus.SUCCESS} + assert statuses[rid2] == {"r1": DownloadStatus.FAILED} + + +class TestConcurrentSameIdCreation: + def test_concurrent_same_id_creation_has_exactly_one_winner(self): + # P1 pin (reproduced by review, then simplified by design): tx_ids are + # attempt-scoped, so racing same-id constructors resolve to exactly ONE + # registered transaction (which owns its outcome slot) and clean ValueErrors + # for every loser -- no interleaving can separate the live transaction from + # its ownership, and no loser leaves any trace. + import threading as th + + service = _make_service() + for round_no in range(30): + tid = f"TX-RACE2-{round_no}" + barrier = th.Barrier(2) + results = [] + + def create(): + barrier.wait() + try: + service.new_transaction(cell=Mock(), timeout=10.0, num_receivers=1, tx_id=tid) + results.append("won") + except ValueError: + results.append("rejected") + + threads = [th.Thread(target=create) for _ in range(2)] + for t in threads: + t.start() + for t in threads: + t.join(5.0) + + assert sorted(results) == ["rejected", "won"], f"round {round_no}: {results}" + live = service._tx_table[tid] + with service._outcome_lock: + owner = service._outcome_owners.get(tid) + assert owner is live, "the live transaction must own its outcome slot" + assert service.get_transaction_outcome(tid) is None + + +class TestShutdownCreationAtomicity: + def test_shutdown_gap_cannot_orphan_a_concurrent_new_transaction(self): + """P-C pin (found by property falsification): shutdown() used to clear _tx_table and + _outcome_owners in SEPARATE critical sections; a new_transaction landing in the gap + registered into both, and the ownership clear then wiped the live transaction's + ownership -- live in _tx_table, outcome unrecordable forever. The teardown is now one + atomic step, so a concurrent creator lands fully before it (torn down) or fully after + it (live AND owning). Deterministic: shutdown's _outcome_lock acquire is held open to + give the racer its best possible window.""" + import threading as th + + service = _make_service() + # warm up: the first new_transaction pays one-time ConfigService resolution for the + # budget defaults -- without this the racer misses the deterministic gap window + warm = service.new_transaction(cell=Mock(), timeout=1.0, num_receivers=1) + service.delete_transaction(warm) + + gap_open = th.Event() + racer_done = th.Event() + real_lock = th.Lock() + + class CoordLock: + """Pauses the shutdown thread's first _outcome_lock acquire until the racer had + its chance -- pre-fix this deterministically reproduced the orphaning.""" + + _tripped = False + + def acquire(self, *a, **k): + if th.current_thread().name == "shutdownT" and not CoordLock._tripped: + CoordLock._tripped = True + # the discriminator is LOCK STATE, not time: post-fix, shutdown holds + # _tx_lock here (atomic teardown) so the racer cannot land in any gap -- + # skip waiting. Pre-fix _tx_lock is free: hold the gap open until the + # racer's new_transaction completes, however slow its cold start is. + if not service._tx_lock.locked(): + gap_open.set() + racer_done.wait(20.0) + return real_lock.acquire(*a, **k) + + def release(self, *a, **k): + return real_lock.release(*a, **k) + + def __enter__(self): + self.acquire() + return self + + def __exit__(self, *exc): + self.release() + + service._outcome_lock = CoordLock() + + def racer(): + # pre-fix: released only when the gap is open. Post-fix the gap never opens; + # the racer proceeds after a short grace and simply lands after the atomic + # teardown (blocked on _tx_lock if it races it). + gap_open.wait(0.5) + service.new_transaction(cell=Mock(), timeout=100.0, num_receivers=1, tx_id="TX-GAP") + racer_done.set() + + shut = th.Thread(target=service.shutdown, name="shutdownT") + race = th.Thread(target=racer) + shut.start() + race.start() + shut.join(10.0) + race.join(10.0) + assert not shut.is_alive() and not race.is_alive() + + # the invariant: whatever the ordering, a tx live in _tx_table owns its outcome slot + live = service._tx_table.get("TX-GAP") + assert live is not None, "the racer's transaction must exist (it ran after teardown)" + with service._outcome_lock: + owner = service._outcome_owners.get("TX-GAP") + assert owner is live, "a live transaction must own its outcome slot even across shutdown" + + # and its outcome is recordable: terminate it and read the verdict + service.delete_transaction("TX-GAP") + assert service.get_transaction_outcome("TX-GAP") is not None + + +class TestConfigResolution: + + def test_budget_config_var_supplies_default(self): + service = _make_service() + with patch.object(ds_module.ConfigService, "get_float_var", return_value=42.0) as gv: + tx_id = service.new_transaction(cell=Mock(), timeout=10.0, num_receivers=1) + tx = service._tx_table[tx_id] + assert tx.receiver_idle_timeout == 42.0 + assert tx.receiver_acquire_timeout == 42.0 + assert gv.call_args.kwargs.get("conf") == ds_module.SystemConfigs.APPLICATION_CONF + + def test_explicit_budget_overrides_config_var(self): + service = _make_service() + with patch.object(ds_module.ConfigService, "get_float_var", return_value=42.0): + tx_id = service.new_transaction(cell=Mock(), timeout=10.0, num_receivers=1, receiver_idle_timeout=7.0) + assert service._tx_table[tx_id].receiver_idle_timeout == 7.0 + + def test_non_positive_config_var_disables_budget(self): + service = _make_service() + with patch.object(ds_module.ConfigService, "get_float_var", return_value=0.0): + tx_id = service.new_transaction(cell=Mock(), timeout=10.0, num_receivers=1) + tx = service._tx_table[tx_id] + assert tx.receiver_idle_timeout is None + assert tx.receiver_acquire_timeout is None + + +class TestTransactionValidation: + def test_receiver_ids_derive_num_receivers(self, caplog): + import logging + + service = _make_service() + with caplog.at_level(logging.WARNING): + tx_id = service.new_transaction(cell=Mock(), timeout=10.0, num_receivers=0, receiver_ids=("a", "b", "a")) + tx = service._tx_table[tx_id] + assert tx.receiver_ids == ("a", "b") # deduped, order kept + assert tx.num_receivers == 2 + # review-requested pin: duplicates are almost certainly a caller bug -- warn + assert any("deduplicated" in r.message for r in caplog.records) + + def test_budget_wider_than_transaction_timeout_warns(self, caplog): + # review-requested pin: a budget >= the transaction timeout can never fire + # (the whole-attempt clock wins) -- dead config must be loud, not silent + import logging + + service = _make_service() + with caplog.at_level(logging.WARNING): + service.new_transaction( + cell=Mock(), timeout=10.0, num_receivers=1, receiver_ids=("r1",), receiver_idle_timeout=10.0 + ) + assert any("can never fire" in r.message for r in caplog.records) + + def test_receiver_ids_num_receivers_mismatch_raises(self): + service = _make_service() + with pytest.raises(ValueError, match="does not match"): + service.new_transaction(cell=Mock(), timeout=10.0, num_receivers=3, receiver_ids=("a", "b")) + + def test_min_receivers_validation(self): + service = _make_service() + with pytest.raises(ValueError, match="must be positive"): + service.new_transaction(cell=Mock(), timeout=10.0, num_receivers=2, min_receivers=0) + with pytest.raises(ValueError, match="exceeds num_receivers"): + service.new_transaction(cell=Mock(), timeout=10.0, num_receivers=2, min_receivers=3) + + def test_negative_budget_rejected(self): + service = _make_service() + with pytest.raises(ValueError, match="must > 0"): + service.new_transaction(cell=Mock(), timeout=10.0, num_receivers=1, receiver_idle_timeout=-1.0) + + +class TestQuorumSurface: + def _refs_outcome(self, statuses, min_receivers, num_receivers=2): + from nvflare.fuel.f3.streaming.transfer_outcome import RefOutcome, TransactionDoneStatus + + refs = [RefOutcome(ref_id="R1", receiver_statuses=statuses)] + return compute_transfer_outcome( + "T1", TransactionDoneStatus.FINISHED, num_receivers, refs, 100.0, min_receivers=min_receivers + ) + + def test_quorum_met_with_partial_fanout(self): + outcome = self._refs_outcome({"a": DownloadStatus.SUCCESS, "b": DownloadStatus.FAILED}, min_receivers=1) + assert not outcome.completed # strict certificate unchanged + assert outcome.quorum_met + + def test_quorum_not_met(self): + outcome = self._refs_outcome({"a": DownloadStatus.SUCCESS, "b": DownloadStatus.FAILED}, min_receivers=2) + assert not outcome.quorum_met + + def test_quorum_falls_back_to_completed_when_unset(self): + success = self._refs_outcome({"a": DownloadStatus.SUCCESS, "b": DownloadStatus.SUCCESS}, min_receivers=None) + assert success.completed and success.quorum_met + partial = self._refs_outcome({"a": DownloadStatus.SUCCESS, "b": DownloadStatus.FAILED}, min_receivers=None) + assert not partial.quorum_met + + def test_quorum_requires_same_receiver_across_all_refs(self): + # disjoint success subsets per ref must NOT satisfy the quorum: no single receiver + # holds the complete (multi-ref) payload + from nvflare.fuel.f3.streaming.transfer_outcome import RefOutcome, TransactionDoneStatus + + refs = [ + RefOutcome(ref_id="R1", receiver_statuses={"a": DownloadStatus.SUCCESS, "b": DownloadStatus.FAILED}), + RefOutcome(ref_id="R2", receiver_statuses={"a": DownloadStatus.FAILED, "b": DownloadStatus.SUCCESS}), + ] + outcome = compute_transfer_outcome("T1", TransactionDoneStatus.FINISHED, 2, refs, 100.0, min_receivers=1) + assert not outcome.quorum_met # each ref has one success, but no common receiver + + refs2 = [ + RefOutcome(ref_id="R1", receiver_statuses={"a": DownloadStatus.SUCCESS, "b": DownloadStatus.FAILED}), + RefOutcome(ref_id="R2", receiver_statuses={"a": DownloadStatus.SUCCESS, "b": DownloadStatus.FAILED}), + ] + outcome2 = compute_transfer_outcome("T1", TransactionDoneStatus.FINISHED, 2, refs2, 100.0, min_receivers=1) + assert outcome2.quorum_met # receiver "a" holds the complete payload + + def test_quorum_fails_closed_with_no_refs(self): + from nvflare.fuel.f3.streaming.transfer_outcome import TransactionDoneStatus + + outcome = compute_transfer_outcome("T1", TransactionDoneStatus.FINISHED, 2, [], 100.0, min_receivers=1) + assert not outcome.quorum_met diff --git a/tests/unit_test/fuel/f3/streaming/receiver_confirm_test.py b/tests/unit_test/fuel/f3/streaming/receiver_confirm_test.py new file mode 100644 index 0000000000..a84abcf30f --- /dev/null +++ b/tests/unit_test/fuel/f3/streaming/receiver_confirm_test.py @@ -0,0 +1,479 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. 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. + +"""Tests for receiver-confirmed completion + retry-aware accounting. + +Producer side: a confirm-capable receiver's served EOF/ERROR is PROVISIONAL; the receiver's +confirmation finalizes it (receiver truth wins, first confirm is final, retries overwrite +provisional state). Legacy receivers keep today's producer-served semantics -- both version +skews degrade to the current behavior, and a runtime kill-switch disables the wire behavior +entirely without a code revert. + +Receiver side: capability is advertised per request, confirmations are sent fire-and-forget +only toward producers that advertised they consume them, and the receiver truth is decided by +Consumer.download_completed() (finalization), not by the served EOF. +""" + +import time +from unittest.mock import Mock, patch + +import pytest + +from nvflare.fuel.f3.cellnet.defs import MessageHeaderKey, ReturnCode +from nvflare.fuel.f3.cellnet.utils import make_reply +from nvflare.fuel.f3.streaming import download_service as ds_module +from nvflare.fuel.f3.streaming.download_service import Consumer, DownloadStatus, ProduceRC, _PropKey, download_object +from tests.unit_test.fuel.f3.streaming.download_test_utils import MockDownloadable +from tests.unit_test.fuel.f3.streaming.download_test_utils import confirm_request as _confirm_request +from tests.unit_test.fuel.f3.streaming.download_test_utils import make_service_no_monitor as _make_service +from tests.unit_test.fuel.f3.streaming.download_test_utils import pull_request as _pull_request +from tests.unit_test.fuel.f3.streaming.download_test_utils import pull_to_terminal as _pull_to_terminal +from tests.unit_test.fuel.f3.streaming.download_test_utils import run_monitor_once, serve_nonce + + +def _new_tx(service, num_receivers=1, timeout=10.0, chunks=1): + tx_id = service.new_transaction(cell=Mock(), timeout=timeout, num_receivers=num_receivers) + obj = MockDownloadable([b"chunk"] * chunks) + rid = service.add_object(tx_id, obj) + return tx_id, rid, obj + + +def _ref(service, rid): + return service._ref_table[rid] + + +class TestProducerSide: + def test_legacy_receiver_finalizes_at_serve(self): + service = _make_service() + tx_id, rid, _ = _new_tx(service) + + reply = _pull_to_terminal(service, rid, "r1", confirm_capable=False) + + assert _PropKey.CONFIRM_EXPECTED not in reply.payload + ref = _ref(service, rid) + assert ref.snapshot_receiver_statuses() == {"r1": DownloadStatus.SUCCESS} + assert ref.snapshot_pending_confirms() == {} + assert service._tx_table[tx_id].is_finished() + + def test_confirm_capable_receiver_is_provisional_at_serve(self): + service = _make_service() + tx_id, rid, _ = _new_tx(service) + + reply = _pull_to_terminal(service, rid, "r1", confirm_capable=True) + + assert reply.payload.get(_PropKey.CONFIRM_EXPECTED) is True + ref = _ref(service, rid) + # served EOF is NOT the receiver's truth yet + assert ref.snapshot_receiver_statuses() == {} + assert ref.snapshot_pending_confirms() == {"r1": DownloadStatus.SUCCESS} + assert not service._tx_table[tx_id].is_finished() + + def test_confirmation_finalizes_and_finishes(self): + service = _make_service() + tx_id, rid, _ = _new_tx(service) + terminal = _pull_to_terminal(service, rid, "r1", confirm_capable=True) + + reply = service._handle_download(_confirm_request(rid, "r1", DownloadStatus.SUCCESS, serve_nonce(terminal))) + + assert reply.get_header(MessageHeaderKey.RETURN_CODE) == ReturnCode.OK + ref = _ref(service, rid) + assert ref.snapshot_receiver_statuses() == {"r1": DownloadStatus.SUCCESS} + assert ref.snapshot_pending_confirms() == {} + assert service._tx_table[tx_id].is_finished() + + def test_receiver_truth_wins_failed_confirm_after_served_eof(self): + # the motivating case: producer served EOF, but the receiver's finalization failed + service = _make_service() + tx_id, rid, _ = _new_tx(service) + terminal = _pull_to_terminal(service, rid, "r1", confirm_capable=True) + + service._handle_download(_confirm_request(rid, "r1", DownloadStatus.FAILED, serve_nonce(terminal))) + + ref = _ref(service, rid) + assert ref.snapshot_receiver_statuses() == {"r1": DownloadStatus.FAILED} + # the transaction reaches its receiver count, but the aggregate outcome must fail + assert service._tx_table[tx_id].is_finished() + run_monitor_once(service, now=time.time()) + outcome = service.get_transaction_outcome(tx_id) + assert not outcome.completed + + def test_retry_overwrites_provisional_and_confirm_wins(self): + # retry-aware accounting: a served ERROR is provisional; a healed retry (EOF) plus a + # SUCCESS confirmation must finalize SUCCESS, not stick at first failure + service = _make_service() + tx_id, rid, obj = _new_tx(service) + ref = _ref(service, rid) + + # simulate a produce-time failure serve, provisionally recorded + ref.obj_served("r1", DownloadStatus.FAILED, expect_confirm=True) + assert ref.snapshot_pending_confirms() == {"r1": DownloadStatus.FAILED} + + # the receiver retries; this time the pull succeeds end-to-end + terminal = _pull_to_terminal(service, rid, "r1", confirm_capable=True) + assert ref.snapshot_pending_confirms() == {"r1": DownloadStatus.SUCCESS} + + service._handle_download(_confirm_request(rid, "r1", DownloadStatus.SUCCESS, serve_nonce(terminal))) + assert ref.snapshot_receiver_statuses() == {"r1": DownloadStatus.SUCCESS} + run_monitor_once(service, now=time.time()) + assert service.get_transaction_outcome(tx_id).completed + + def test_first_confirm_is_final(self): + service = _make_service() + tx_id, rid, _ = _new_tx(service) + terminal = _pull_to_terminal(service, rid, "r1", confirm_capable=True) + + service._handle_download(_confirm_request(rid, "r1", DownloadStatus.SUCCESS, serve_nonce(terminal))) + service._handle_download(_confirm_request(rid, "r1", DownloadStatus.FAILED, serve_nonce(terminal))) + + assert _ref(service, rid).snapshot_receiver_statuses() == {"r1": DownloadStatus.SUCCESS} + + def test_late_duplicate_serve_cannot_resurrect_provisional(self): + service = _make_service() + tx_id, rid, _ = _new_tx(service) + terminal = _pull_to_terminal(service, rid, "r1", confirm_capable=True) + service._handle_download(_confirm_request(rid, "r1", DownloadStatus.SUCCESS, serve_nonce(terminal))) + + ref = _ref(service, rid) + ref.obj_served("r1", DownloadStatus.FAILED, expect_confirm=True) + + assert ref.snapshot_pending_confirms() == {} + assert ref.snapshot_receiver_statuses() == {"r1": DownloadStatus.SUCCESS} + + def test_invalid_confirm_status_is_ignored(self): + service = _make_service() + tx_id, rid, _ = _new_tx(service) + _pull_to_terminal(service, rid, "r1", confirm_capable=True) + + service._handle_download(_confirm_request(rid, "r1", "bogus")) + + ref = _ref(service, rid) + assert ref.snapshot_receiver_statuses() == {} + assert ref.snapshot_pending_confirms() == {"r1": DownloadStatus.SUCCESS} + + def test_unsolicited_confirm_cannot_certify(self): + # fail-open hole guard: a CONFIRM for a receiver that was never served a terminal + # reply in this life of the ref must be dropped -- otherwise a stale confirm + # delayed across a ref_id reuse could certify (or pre-poison) a transfer that never + # delivered a byte + service = _make_service() + tx_id, rid, _ = _new_tx(service) + + service._handle_download(_confirm_request(rid, "r1", DownloadStatus.SUCCESS)) + + ref = _ref(service, rid) + assert ref.snapshot_receiver_statuses() == {} + assert not service._tx_table[tx_id].is_finished() + run_monitor_once(service, now=time.time()) + assert service.get_transaction_outcome(tx_id) is None # still live, nothing certified + + # the FAILED mirror cannot pre-poison either + service._handle_download(_confirm_request(rid, "r1", DownloadStatus.FAILED)) + assert ref.snapshot_receiver_statuses() == {} + + def test_tombstoned_ref_replay_does_not_solicit_confirm(self): + # after cleanup, a confirm-capable receiver retrying a finished ref gets the EOF + # replay WITHOUT CONFIRM_EXPECTED (no live ref to confirm against), and a stray + # confirm for the tombstoned rid is dropped OK without disturbing the outcome + service = _make_service() + tx_id, rid, _ = _new_tx(service) + terminal = _pull_to_terminal(service, rid, "r1", confirm_capable=True) + service._handle_download(_confirm_request(rid, "r1", DownloadStatus.SUCCESS, serve_nonce(terminal))) + run_monitor_once(service, now=time.time()) # FINISHED -> tombstoned + assert rid not in service._ref_table + + replay = service._handle_download(_pull_request(rid, "r1", confirm_capable=True)) + assert replay.payload.get(_PropKey.STATUS) == ProduceRC.EOF + assert _PropKey.CONFIRM_EXPECTED not in replay.payload + + reply = service._handle_download(_confirm_request(rid, "r1", DownloadStatus.FAILED)) + assert reply.get_header(MessageHeaderKey.RETURN_CODE) == ReturnCode.OK + assert service.get_transaction_outcome(tx_id).completed # unchanged + + def test_kill_switch_reads_application_conf(self): + # the switch resolves through the standard application-config source (env var + # NVFLARE_STREAMING_RECEIVER_CONFIRM_ENABLED / job config), like neighboring vars + with patch.object(ds_module, "_receiver_confirm_cached", None): + with patch.object(ds_module.ConfigService, "get_bool_var", return_value=False) as gv: + assert ds_module._receiver_confirm_enabled() is False + assert gv.call_args.kwargs.get("conf") == ds_module.SystemConfigs.APPLICATION_CONF + + def test_unexpected_receiver_cannot_complete_declared_identities(self): + # P1 pin: with receiver_ids=("a", "b"), a success from "a" plus a success from an + # UNEXPECTED receiver "x" must not certify -- "b" never received anything + service = _make_service() + tx_id = service.new_transaction(cell=Mock(), timeout=10.0, num_receivers=0, receiver_ids=("a", "b")) + obj = MockDownloadable([b"chunk"]) + rid = service.add_object(tx_id, obj) + + _pull_to_terminal(service, rid, "a") # expected, succeeds + _pull_to_terminal(service, rid, "x") # unexpected, also succeeds + + tx = service._tx_table[tx_id] + assert not tx.is_finished(), "count-based completion would wrongly finish here" + assert not obj.released, "downloaded_to_all/release must not fire without 'b'" + + _pull_to_terminal(service, rid, "b") # the missing expected receiver arrives + assert tx.is_finished() + run_monitor_once(service, now=time.time()) + outcome = service.get_transaction_outcome(tx_id) + assert outcome.completed + assert outcome.receiver_ids == ("a", "b") + + def test_declared_identity_missing_fails_outcome_despite_count(self): + # two successes are recorded ("a" + unexpected "x"), matching the receiver COUNT -- + # but declared receiver "b" got nothing, so the outcome must fail closed (here via + # the transaction TTL, since "b" never pulled and no acquire budget was set) + service = _make_service() + tx_id = service.new_transaction( + cell=Mock(), timeout=10.0, num_receivers=0, receiver_ids=("a", "b"), receiver_idle_timeout=5.0 + ) + rid = service.add_object(tx_id, MockDownloadable([b"chunk"])) + _pull_to_terminal(service, rid, "a") + _pull_to_terminal(service, rid, "x") # unexpected + + run_monitor_once(service, now=time.time() + 30.0) + outcome = service.get_transaction_outcome(tx_id) + assert outcome is not None + assert not outcome.completed, "two successes without declared receiver 'b' must never certify" + assert not outcome.quorum_met or outcome.min_receivers is None + + def test_stale_confirm_from_previous_ref_life_is_dropped_even_with_new_pending(self): + # P1 pin (reproduced by review): old life's confirmation must not finalize the NEW + # life of a reused ref_id even when the new life has its own pending serve for the + # same receiver -- the per-serve nonce is what tells the lives apart + service = _make_service() + tx1 = service.new_transaction(cell=Mock(), timeout=10.0, num_receivers=1, tx_id="TX-LIFE") + rid = service.add_object(tx1, MockDownloadable([b"chunk"]), ref_id="R-LIFE") + old_terminal = _pull_to_terminal(service, rid, "r1", confirm_capable=True) + old_nonce = serve_nonce(old_terminal) + + # the transfer is retried as a NEW attempt (tx_ids are attempt-scoped, never + # reused) that re-serves the SAME ref_id -- the real pattern is PASS_THROUGH + # re-emission, where the ref_id was minted before registration and re-added. + # First retire the old attempt so the rid can be re-registered. + service.delete_transaction(tx1) + tx2 = service.new_transaction(cell=Mock(), timeout=10.0, num_receivers=1, tx_id="TX-LIFE-2") + rid2 = service.add_object(tx2, MockDownloadable([b"chunk"]), ref_id="R-LIFE") + assert rid2 == rid + new_terminal = _pull_to_terminal(service, rid, "r1", confirm_capable=True) # new pending serve exists + + # the OLD life's delayed confirmation arrives: wrong nonce -> dropped + service._handle_download(_confirm_request(rid, "r1", DownloadStatus.SUCCESS, old_nonce)) + ref = _ref(service, rid) + assert ref.snapshot_receiver_statuses() == {} + assert ref.snapshot_pending_confirms() == {"r1": DownloadStatus.SUCCESS} + + # the NEW life's confirmation (correct nonce) finalizes + service._handle_download(_confirm_request(rid, "r1", DownloadStatus.SUCCESS, serve_nonce(new_terminal))) + assert ref.snapshot_receiver_statuses() == {"r1": DownloadStatus.SUCCESS} + + def test_confirm_message_dispatches_to_confirm_handling_not_pull(self): + # review-requested pin: a message carrying the CONFIRM key must route to + # _handle_confirm and never fall through to pull handling -- a fall-through + # would treat the confirm as an initial pull request (state=None) and + # re-serve the object from the beginning + service = _make_service() + tx1 = service.new_transaction(cell=Mock(), timeout=10.0, num_receivers=1) + obj = MockDownloadable([b"chunk"]) + rid = service.add_object(tx1, obj, ref_id="R-DISPATCH") + + reply = service._handle_download(_confirm_request(rid, "r1", DownloadStatus.SUCCESS, "bogus-nonce")) + + # confirm-path reply: bare OK, no chunk payload -- and produce() never ran + assert reply.payload is None or _PropKey.DATA not in (reply.payload or {}) + assert obj.current_chunk == 0 + ref = _ref(service, rid) + # the unsolicited confirm was dropped by confirm handling (no pending serve) + assert ref.snapshot_receiver_statuses() == {} + assert ref.snapshot_pending_confirms() == {} + + def test_late_confirm_for_unknown_ref_is_dropped_ok(self): + service = _make_service() + reply = service._handle_download(_confirm_request("R-GONE", "r1", DownloadStatus.SUCCESS)) + assert reply.get_header(MessageHeaderKey.RETURN_CODE) == ReturnCode.OK + + def test_unconfirmed_receiver_fails_closed_at_timeout(self): + # a lost fire-and-forget confirmation must not certify success: at the transaction + # timeout the unconfirmed receiver has no final status and the outcome fails + service = _make_service() + tx_id, rid, _ = _new_tx(service, timeout=10.0) + _pull_to_terminal(service, rid, "r1", confirm_capable=True) + + tx = service._tx_table[tx_id] + tx.last_active_time = time.time() - 11.0 + run_monitor_once(service, now=time.time()) + + outcome = service.get_transaction_outcome(tx_id) + assert outcome is not None + assert not outcome.completed + + def test_kill_switch_off_restores_legacy_semantics(self): + service = _make_service() + with patch.object(ds_module, "_receiver_confirm_cached", False): + tx_id, rid, _ = _new_tx(service) + reply = _pull_to_terminal(service, rid, "r1", confirm_capable=True) + + # producer ignores the advertised capability entirely + assert _PropKey.CONFIRM_EXPECTED not in reply.payload + ref = _ref(service, rid) + assert ref.snapshot_receiver_statuses() == {"r1": DownloadStatus.SUCCESS} + assert ref.snapshot_pending_confirms() == {} + assert service._tx_table[tx_id].is_finished() + + def test_mixed_fleet_finishes_on_mixed_semantics(self): + # one legacy receiver (final at serve) + one confirm-capable receiver (final at confirm) + service = _make_service() + tx_id, rid, _ = _new_tx(service, num_receivers=2) + tx = service._tx_table[tx_id] + + _pull_to_terminal(service, rid, "legacy", confirm_capable=False) + terminal = _pull_to_terminal(service, rid, "modern", confirm_capable=True) + assert not tx.is_finished() # modern receiver not confirmed yet + + service._handle_download(_confirm_request(rid, "modern", DownloadStatus.SUCCESS, serve_nonce(terminal))) + assert tx.is_finished() + run_monitor_once(service, now=time.time()) + assert service.get_transaction_outcome(tx_id).completed + + +class _ScriptedCell: + """A cell whose send_request returns scripted replies and which records fire_and_forget.""" + + def __init__(self, replies): + self._replies = list(replies) + self.requests = [] + self.confirms = [] + self.confirm_kwargs = [] + + def send_request(self, channel, target, topic, request, timeout, secure, optional, abort_signal): + self.requests.append(request.payload) + if not self._replies: + raise AssertionError("scripted cell ran out of replies") + return self._replies.pop(0) + + def fire_and_forget(self, channel, topic, targets, message, secure=False, optional=False): + self.confirms.append(message.payload) + self.confirm_kwargs.append({"secure": secure, "optional": optional}) + + +class _RecordingConsumer(Consumer): + def __init__(self, fail_on_complete=False): + super().__init__() + self.consumed = [] + self.completed = False + self.failed_reason = None + self._fail_on_complete = fail_on_complete + + def consume(self, ref_id, state, data): + self.consumed.append(data) + return state or {} + + def download_completed(self, ref_id): + if self._fail_on_complete: + raise RuntimeError("finalization failed") + self.completed = True + + def download_failed(self, ref_id, reason): + self.failed_reason = reason + + +def _ok_reply(body): + return make_reply(ReturnCode.OK, body=body) + + +def _chunk_reply(confirm_expected=True): + body = {_PropKey.STATUS: ProduceRC.OK, _PropKey.STATE: {"seq": 1}, _PropKey.DATA: b"chunk"} + if confirm_expected: + body[_PropKey.CONFIRM_EXPECTED] = True + return _ok_reply(body) + + +def _terminal_reply(status, confirm_expected=True, nonce="n-test"): + body = {_PropKey.STATUS: status} + if confirm_expected: + body[_PropKey.CONFIRM_EXPECTED] = True + body[_PropKey.CONFIRM_NONCE] = nonce + return _ok_reply(body) + + +class TestReceiverSide: + def test_advertises_capability_and_confirms_success_after_finalization(self): + cell = _ScriptedCell([_chunk_reply(), _terminal_reply(ProduceRC.EOF)]) + consumer = _RecordingConsumer() + + download_object(from_fqcn="site-1", ref_id="R1", per_request_timeout=5.0, cell=cell, consumer=consumer) + + assert all(req.get(_PropKey.CONFIRM_CAPABLE) is True for req in cell.requests) + assert consumer.completed + assert cell.confirms == [ + {_PropKey.REF_ID: "R1", _PropKey.CONFIRM: DownloadStatus.SUCCESS, _PropKey.CONFIRM_NONCE: "n-test"} + ] + + def test_confirms_failed_when_finalization_raises(self): + # served EOF but download_completed raises: the producer must learn receiver truth + cell = _ScriptedCell([_terminal_reply(ProduceRC.EOF)]) + consumer = _RecordingConsumer(fail_on_complete=True) + + with pytest.raises(RuntimeError, match="finalization failed"): + download_object(from_fqcn="site-1", ref_id="R1", per_request_timeout=5.0, cell=cell, consumer=consumer) + + assert cell.confirms == [ + {_PropKey.REF_ID: "R1", _PropKey.CONFIRM: DownloadStatus.FAILED, _PropKey.CONFIRM_NONCE: "n-test"} + ] + + def test_confirms_failed_on_producer_error(self): + cell = _ScriptedCell([_chunk_reply(), _terminal_reply(ProduceRC.ERROR)]) + consumer = _RecordingConsumer() + + download_object(from_fqcn="site-1", ref_id="R1", per_request_timeout=5.0, cell=cell, consumer=consumer) + + assert consumer.failed_reason is not None + assert cell.confirms == [ + {_PropKey.REF_ID: "R1", _PropKey.CONFIRM: DownloadStatus.FAILED, _PropKey.CONFIRM_NONCE: "n-test"} + ] + + def test_no_confirm_toward_legacy_producer(self): + # the producer never advertised CONFIRM_EXPECTED: a new receiver must send nothing extra + cell = _ScriptedCell( + [_chunk_reply(confirm_expected=False), _terminal_reply(ProduceRC.EOF, confirm_expected=False)] + ) + consumer = _RecordingConsumer() + + download_object(from_fqcn="site-1", ref_id="R1", per_request_timeout=5.0, cell=cell, consumer=consumer) + + assert consumer.completed + assert cell.confirms == [] + + def test_confirm_honors_secure_flag(self): + # P2 pin: the confirmation must ride with the same security posture as the + # download requests it concludes + cell = _ScriptedCell([_terminal_reply(ProduceRC.EOF)]) + consumer = _RecordingConsumer() + + download_object( + from_fqcn="site-1", ref_id="R1", per_request_timeout=5.0, cell=cell, consumer=consumer, secure=True + ) + + assert cell.confirm_kwargs == [{"secure": True, "optional": False}] + + def test_kill_switch_off_receiver_is_fully_legacy(self): + with patch.object(ds_module, "_receiver_confirm_cached", False): + cell = _ScriptedCell([_terminal_reply(ProduceRC.EOF, confirm_expected=True)]) + consumer = _RecordingConsumer() + + download_object(from_fqcn="site-1", ref_id="R1", per_request_timeout=5.0, cell=cell, consumer=consumer) + + assert all(_PropKey.CONFIRM_CAPABLE not in req for req in cell.requests) + assert cell.confirms == [] + assert consumer.completed diff --git a/tests/unit_test/fuel/f3/streaming/transfer_outcome_test.py b/tests/unit_test/fuel/f3/streaming/transfer_outcome_test.py index 111c0eaf6e..41459aac39 100644 --- a/tests/unit_test/fuel/f3/streaming/transfer_outcome_test.py +++ b/tests/unit_test/fuel/f3/streaming/transfer_outcome_test.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import dataclasses +import threading import time import pytest @@ -35,6 +36,30 @@ def _stub_obj(): return MockDownloadable([b"chunk"]) +class TestOutcomeImmutability: + """The recorded outcome is shared with pollers and outcome_cb consumers: it must be deep-frozen.""" + + def test_outcome_is_deep_frozen(self): + source_statuses = {"r1": DownloadStatus.SUCCESS} + ref = RefOutcome(ref_id="R1", receiver_statuses=source_statuses) + outcome = compute_transfer_outcome("T1", TransactionDoneStatus.FINISHED, 1, [ref], 100.0) + + # containers are frozen, not just the dataclass attributes + assert isinstance(outcome.refs, tuple) + with pytest.raises(TypeError): + outcome.refs[0].receiver_statuses["r2"] = DownloadStatus.SUCCESS + with pytest.raises(dataclasses.FrozenInstanceError): + outcome.refs = () + with pytest.raises(dataclasses.FrozenInstanceError): + outcome.refs[0].receiver_statuses = {} + + # the frozen view is a private copy: mutating the source dict after + # construction cannot rewrite the recorded per-receiver truth + source_statuses["r1"] = DownloadStatus.FAILED + assert outcome.refs[0].receiver_statuses == {"r1": DownloadStatus.SUCCESS} + assert outcome.completed + + class TestComputeTransferOutcome: """Aggregation rules: COMPLETED only when every expected receiver succeeded; receiver truth wins.""" @@ -151,7 +176,35 @@ def test_transaction_done_returns_outcome_with_receiver_map(self): assert outcome.refs[0].receiver_statuses == {"r1": DownloadStatus.SUCCESS} assert obj.released - def test_on_outcome_fires_before_user_callbacks(self): + def test_raising_release_cannot_leave_waiter_unresolved(self): + """P1 pin: an unguarded custom release() used to escape transaction_done -- with + recording moved to the end, that left the outcome unrecorded, the waiter pending + forever, ownership registered, and (from the monitor) killed the monitor thread. + Releases are individually guarded and recording runs in a finally.""" + from unittest.mock import Mock + + service = make_isolated_download_service() + service._tx_monitor = Mock() + cell = Mock() + tx_id = service.new_transaction(cell=cell, timeout=10.0, num_receivers=1) + obj = _stub_obj() + obj.release = Mock(side_effect=RuntimeError("release blew up")) + service.add_object(tx_id, obj) + ref = service._tx_table[tx_id].snapshot_refs()[0] + ref.obj_downloaded("r1", DownloadStatus.SUCCESS) + waiter = service.get_transfer_waiter(tx_id) + + run_monitor_once(service, now=time.time()) # must not raise off the monitor thread + + outcome = waiter.wait(timeout=5.0) + assert outcome is not None and outcome.completed, "recording must survive a raising release()" + with service._outcome_lock: + assert tx_id not in service._outcome_owners, "ownership must be consumed" + + def test_on_outcome_fires_after_callbacks_and_release(self): + # settle-then-record: recording (which releases waiters) happens only after the + # callback chain and source release complete, so an upper layer acting on + # waiter.wait() can never preempt them (e.g. by stopping the producer process) order = [] tx = _Transaction( timeout=10.0, @@ -160,12 +213,13 @@ def test_on_outcome_fires_before_user_callbacks(self): outcome_cb=lambda outcome: order.append("outcome_cb"), ) obj = _stub_obj() + obj.release = lambda: order.append("release") # observe source release order ref = tx.add_object(obj) ref.obj_downloaded("r1", DownloadStatus.FAILED) tx.transaction_done(TransactionDoneStatus.FINISHED, on_outcome=lambda outcome: order.append("recorded")) - assert order == ["recorded", "done_cb", "outcome_cb"] + assert order == ["done_cb", "outcome_cb", "release", "recorded"] def test_raising_callbacks_do_not_break_recording_or_release(self): # a raising transaction_done_cb must not skip outcome recording, source @@ -216,11 +270,11 @@ def _add_tx(self, service, num_receivers=1, timeout=10.0): tx = _Transaction(timeout=timeout, num_receivers=num_receivers) with service._tx_lock: service._tx_table[tx.tid] = tx - # mirror new_transaction(): a live tx is registered as the current incarnation. - # _record_outcome() records only for the live incarnation (current is tx), so a tx - # placed directly in _tx_table without this would never record its terminal outcome. + # mirror new_transaction(): a live tx owns its outcome slot. + # _record_outcome() records only for the owning tx, so a tx placed directly in + # _tx_table without ownership would never record its terminal outcome. with service._outcome_lock: - service._tx_incarnations[tx.tid] = tx + service._outcome_owners[tx.tid] = tx obj = _stub_obj() rid = service.add_object(tx.tid, obj) with service._tx_lock: @@ -285,52 +339,641 @@ def test_delete_after_success_records_completed(self): assert outcome.completed assert outcome.done_status == TransactionDoneStatus.DELETED - def test_reused_tx_id_does_not_surface_stale_outcome(self): - # retry with the same explicit tx_id (design: tx_id = transfer_id, retries reuse it) + def test_new_transaction_rejects_live_id(self): + """tx_ids are attempt-scoped: registering an id that names a LIVE transaction + raises immediately -- and, unlike the old retire-on-reuse design, the live + transaction is completely undisturbed (no callbacks, no release, still servable).""" + from unittest.mock import Mock + + service = make_isolated_download_service() + service._tx_monitor = Mock() + cell = Mock() + done, outcome_cbs = [], [] + service.new_transaction( + cell=cell, + timeout=10.0, + num_receivers=1, + tx_id="TX-DUP", + transaction_done_cb=lambda tid, status, base_objs, **kw: done.append(status), + outcome_cb=lambda outcome: outcome_cbs.append(outcome), + ) + obj = _stub_obj() + rid = service.add_object("TX-DUP", obj) + + with pytest.raises(ValueError, match="already in use"): + service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-DUP") + + # the duplicate attempt left no trace on the live transaction + assert rid in service._ref_table + assert not obj.released + assert done == [] and outcome_cbs == [] + with service._tx_lock: + with service._outcome_lock: + assert service._tx_table["TX-DUP"] is service._outcome_owners["TX-DUP"] + + def test_new_transaction_rejects_settling_and_receipted_id(self): + """The duplicate check covers the id's whole lifetime: live (above), SETTLING + (popped from _tx_table, ownership not yet consumed), and RECEIPTED (outcome + retained for TX_OUTCOME_TTL). Only receipt expiry frees a used id.""" + from functools import partial + from unittest.mock import Mock + + service = make_isolated_download_service() + service._tx_monitor = Mock() + cell = Mock() + service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-LIFE") + with service._tx_lock: + old_tx = service._tx_table.pop("TX-LIFE") # monitor termination step 1 + + # mid-settlement: ownership still registered + with pytest.raises(ValueError, match="already in use"): + service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-LIFE") + + old_tx.transaction_done(TransactionDoneStatus.FINISHED, on_outcome=partial(service._record_outcome, tx=old_tx)) + + # receipted: the verdict is retained and the id stays excluded + assert service.get_transaction_outcome("TX-LIFE") is not None + with pytest.raises(ValueError, match="already in use"): + service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-LIFE") + + # receipt expiry (TTL) frees the id -- checked INLINE by new_transaction, so an + # expired-but-unswept receipt does not stretch the exclusion window + recorded = service.get_transaction_outcome("TX-LIFE") + assert recorded is not None + with service._outcome_lock: + service._tx_outcomes["TX-LIFE"] = dataclasses.replace( + recorded, timestamp=recorded.timestamp - service.TX_OUTCOME_TTL - 1 + ) + assert service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-LIFE") == "TX-LIFE" + + def test_drain_leak_poisons_id_until_operation_exits(self): + """P1 pin: a drain-timeout leak (operation hung past OP_DRAIN_TIMEOUT) may resume + and emit later -- so its tx_id must stay excluded from registration even after + the receipt expires, until every leaked operation has exited. Otherwise the + leaked emission could land under a recycled id.""" + from unittest.mock import Mock + + import nvflare.fuel.f3.streaming.download_service as ds_module + + service = make_isolated_download_service() + service._tx_monitor = Mock() + cell = Mock() + service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-LEAK") + with service._tx_lock: + tx = service._tx_table["TX-LEAK"] + assert tx.begin_op() # an in-flight operation (serve/budget pass) that will hang + + saved = ds_module.OP_DRAIN_TIMEOUT + ds_module.OP_DRAIN_TIMEOUT = 0.2 + try: + service.delete_transaction("TX-LEAK") # drain times out; settles anyway (logged) + finally: + ds_module.OP_DRAIN_TIMEOUT = saved + + # force-expire the receipt: the id must STILL be excluded by the leak + recorded = service.get_transaction_outcome("TX-LEAK") + assert recorded is not None + with service._outcome_lock: + service._tx_outcomes["TX-LEAK"] = dataclasses.replace( + recorded, timestamp=recorded.timestamp - service.TX_OUTCOME_TTL - 1 + ) + with pytest.raises(ValueError, match="has not fully terminated"): + service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-LEAK") + + # the leaked operation exits: the id becomes usable + tx.end_op() + assert service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-LEAK") == "TX-LEAK" + + def test_settlement_callback_can_create_transactions_but_not_its_own_id(self): + """Settlement runs outside all service locks, so callbacks may freely call back + in and create NEW transactions; recreating the settling transaction's own id is + rejected like any other in-use id (its ownership is consumed only after the + callbacks, in on_outcome).""" + from unittest.mock import Mock + + service = make_isolated_download_service() + service._tx_monitor = Mock() + cell = Mock() + errors, created = [], [] + + def cb(outcome): + try: + service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-SELF") + except ValueError as e: + errors.append(str(e)) + created.append(service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-OTHER")) + + service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-SELF", outcome_cb=cb) + service.delete_transaction("TX-SELF") + + assert len(errors) == 1 and "has not fully terminated" in errors[0] + assert created == ["TX-OTHER"] + with service._tx_lock: + assert "TX-OTHER" in service._tx_table + + def test_recording_failure_leaves_id_unusable_and_shutdown_frees_waiters(self): + """If outcome recording itself fails abnormally, ownership is never consumed: + the id stays excluded (fail-safe -- fresh uuids never collide with it) and its + waiter is released by shutdown, the terminal never-hang backstop.""" from unittest.mock import Mock service = make_isolated_download_service() - service._tx_monitor = Mock() # suppress real monitor thread start + service._tx_monitor = Mock() cell = Mock() + service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-DEAD") + stranded = service.get_transfer_waiter("TX-DEAD") + + def exploding_record(outcome, tx): + raise RuntimeError("recording blew up") + + service._record_outcome = exploding_record try: - first = service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-REUSE") - service.delete_transaction(first) - assert service.get_transaction_outcome("TX-REUSE") is not None - - second = service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-REUSE") - assert second == "TX-REUSE" - # the stale terminal outcome of the previous incarnation is purged - assert service.get_transaction_outcome("TX-REUSE") is None + service.delete_transaction("TX-DEAD") # settles; recording fails; owner never consumed finally: - service.shutdown() + del service._record_outcome # restore the inherited classmethod + assert not stranded.done() + with pytest.raises(ValueError, match="already in use"): + service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-DEAD") - def test_stale_incarnation_cannot_record_over_live_retry(self): - # the record-after-purge race: termination removes the old tx from _tx_table, - # a retry registers the same tx_id, THEN the old incarnation records its - # outcome — it must not shadow the live retry - from functools import partial + service.shutdown() + assert stranded.done() and stranded.outcome is None + + def test_shutdown_window_keeps_leaked_id_excluded(self): + """P1 pin: shutdown clears the ownership and receipt exclusions up front, and + leak registration used to run only AFTER the deferred settlement -- so a + drain-leaked id looked free for that whole window and a replacement could + register while the leaked operation was still able to emit. Shutdown now + pre-registers every owner as a potential leak under the locks.""" from unittest.mock import Mock + import nvflare.fuel.f3.streaming.download_service as ds_module + service = make_isolated_download_service() - service._tx_monitor = Mock() # suppress real monitor thread start + service._tx_monitor = Mock() cell = Mock() + cb_entered = threading.Event() + cb_release = threading.Event() + + def gated_done_cb(tid, status, base_objs, **kw): + cb_entered.set() + cb_release.wait(10.0) + + service.new_transaction( + cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-SHUTLEAK", transaction_done_cb=gated_done_cb + ) + with service._tx_lock: + tx = service._tx_table["TX-SHUTLEAK"] + assert tx.begin_op() # an in-flight operation that will outlive the drain + + saved = ds_module.OP_DRAIN_TIMEOUT + ds_module.OP_DRAIN_TIMEOUT = 0.2 + shutter = threading.Thread(target=service.shutdown) try: - service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-RACE") + shutter.start() + assert cb_entered.wait(5.0) # drain timed out; deferred settlement mid-window + + # the leaked operation could still emit: its id must NOT look free + with pytest.raises(ValueError, match="has not fully terminated"): + service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-SHUTLEAK") + finally: + cb_release.set() + shutter.join(5.0) + ds_module.OP_DRAIN_TIMEOUT = saved + assert not shutter.is_alive() + + # still excluded after shutdown while the operation is in flight + with pytest.raises(ValueError, match="has not fully terminated"): + service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-SHUTLEAK") + + # the operation exits: the id frees + tx.end_op() + assert service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-SHUTLEAK") == "TX-SHUTLEAK" + + def test_outcome_computation_failure_still_records_fail_closed(self): + """Greptile pin: if compute_transfer_outcome itself raises, settlement used to + skip recording entirely -- ownership stayed consumed-never, and a waiter parked + on the id hung until shutdown. The finally now records a fail-closed fallback + verdict (empty refs certify nothing), so waiters always resolve.""" + from unittest.mock import Mock, patch + + import nvflare.fuel.f3.streaming.download_service as ds_module + + service = make_isolated_download_service() + service._tx_monitor = Mock() + cell = Mock() + obj = _stub_obj() + outcome_cbs = [] + released_when_outcome_cb_ran = [] + + def capture_outcome(outcome): + outcome_cbs.append(outcome) + released_when_outcome_cb_ran.append(obj.released) + + service.new_transaction( + cell=cell, + timeout=10.0, + num_receivers=1, + tx_id="TX-COMPUTE", + receiver_ids=("r1",), + min_receivers=1, + outcome_cb=capture_outcome, + ) + service.add_object("TX-COMPUTE", obj) + waiters = [service.get_transfer_waiter("TX-COMPUTE") for _ in range(3)] + + def always_exploding(*args, **kwargs): + raise RuntimeError("outcome computation blew up") + + # PERSISTENT failure (review round 10: a one-shot failure masks a fallback + # that reuses the same failing function) + with patch.object(ds_module, "compute_transfer_outcome", side_effect=always_exploding): + service.delete_transaction("TX-COMPUTE") + + outcomes = [waiter.wait(timeout=5.0) for waiter in waiters] + outcome = outcomes[0] + assert outcome is not None, "waiters must resolve even when outcome computation raises persistently" + assert all(o is outcome for o in outcomes), "all parked waiters must receive the same recorded receipt" + late_waiter = service.get_transfer_waiter("TX-COMPUTE") + assert late_waiter.done() and late_waiter.wait(timeout=0) is outcome + assert not outcome.completed, "the fallback verdict must be fail-closed" + with service._outcome_lock: + assert "TX-COMPUTE" not in service._outcome_owners, "ownership must be consumed" + # cleanup surfaces still ran: sources freed, per-object hook fired + assert obj.released, "source release must run even when the verdict computation fails" + assert obj.transaction_done_calls, "obj.transaction_done must run even when the verdict computation fails" + # codex P2/P3: the outcome_cb contract holds on the fail-closed path, with + # metadata preserved and an honest reason + assert len(outcome_cbs) == 1 + assert released_when_outcome_cb_ran == [False], "outcome_cb must run before source release" + assert outcome_cbs[0].reason == TransferOutcomeReason.COMPUTATION_FAILED + assert outcome_cbs[0].receiver_ids == ("r1",) + assert outcome_cbs[0].min_receivers == 1 + assert outcome.reason == TransferOutcomeReason.COMPUTATION_FAILED + assert outcome.receiver_ids == ("r1",) + assert outcome.min_receivers == 1 + # TX-COMPUTE's own termination marker is released (same-id re-registration + # stays blocked by the retained receipt, which is intentional) + with service._tx_lock: + assert "TX-COMPUTE" not in service._terminating_txs + + @pytest.mark.parametrize( + "terminal_path,expected_done_status", + [ + ("finished", TransactionDoneStatus.FINISHED), + ("timeout", TransactionDoneStatus.TIMEOUT), + ], + ) + def test_monitor_terminal_paths_survive_persistent_outcome_computation_failure( + self, terminal_path, expected_done_status + ): + """A bad verdict computation must not kill the real monitor termination paths. + + FINISHED and TIMEOUT are classified by the monitor rather than by an explicit + delete call, so both must still release sources, record a fail-closed receipt, + resolve waiters, and retire ownership/termination markers. + """ + from unittest.mock import Mock, patch + + import nvflare.fuel.f3.streaming.download_service as ds_module + + service = make_isolated_download_service() + service._tx_monitor = Mock() + cell = Mock() + tx_id = service.new_transaction(cell=cell, timeout=10.0, num_receivers=1) + obj = _stub_obj() + rid = service.add_object(tx_id, obj) + if terminal_path == "finished": + with service._tx_lock: + ref = service._ref_table[rid] + ref.obj_downloaded("r1", DownloadStatus.SUCCESS) + monitor_now = time.time() + else: with service._tx_lock: - old_tx = service._tx_table.pop("TX-RACE") # termination step 1, as the monitor does + tx = service._tx_table[tx_id] + monitor_now = tx.last_active_time + tx.timeout + 1.0 + waiter = service.get_transfer_waiter(tx_id) + + with patch.object(ds_module, "compute_transfer_outcome", side_effect=RuntimeError("persistent failure")): + run_monitor_once(service, now=monitor_now) + + outcome = waiter.wait(timeout=5.0) + assert outcome is not None + assert outcome.status == TransferProgressState.FAILED + assert outcome.reason == TransferOutcomeReason.COMPUTATION_FAILED + assert outcome.done_status == expected_done_status + assert obj.released + assert obj.transaction_done_calls == [(tx_id, expected_done_status)] + with service._outcome_lock: + assert tx_id not in service._outcome_owners + with service._tx_lock: + assert tx_id not in service._terminating_txs + + def test_computation_failure_with_drain_leak_keeps_id_excluded_until_operation_exits(self): + """The fail-closed settlement and drain-leak protections must compose: the + waiter resolves even though an operation missed the drain deadline, but the old + tx_id remains excluded after receipt expiry until that operation really exits. + """ + from unittest.mock import Mock, patch + + import nvflare.fuel.f3.streaming.download_service as ds_module - # the retry registers the same id before the old incarnation records - service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-RACE") + service = make_isolated_download_service() + service._tx_monitor = Mock() + cell = Mock() + tx_id = "TX-COMPUTE-LEAK" + service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id=tx_id) + service.add_object(tx_id, _stub_obj()) + with service._tx_lock: + tx = service._tx_table[tx_id] + assert tx.begin_op() + waiter = service.get_transfer_waiter(tx_id) + + saved = ds_module.OP_DRAIN_TIMEOUT + ds_module.OP_DRAIN_TIMEOUT = 0.05 + try: + with patch.object(ds_module, "compute_transfer_outcome", side_effect=RuntimeError("persistent failure")): + service.delete_transaction(tx_id) + finally: + ds_module.OP_DRAIN_TIMEOUT = saved - # the old incarnation now finishes terminating and tries to record - old_tx.transaction_done( - TransactionDoneStatus.DELETED, on_outcome=partial(service._record_outcome, tx=old_tx) + outcome = waiter.wait(timeout=5.0) + assert outcome is not None and outcome.reason == TransferOutcomeReason.COMPUTATION_FAILED + with service._tx_lock: + assert service._terminating_txs.get(tx_id) is tx + + # Receipt expiry alone is insufficient while an old operation can still emit. + with service._outcome_lock: + service._tx_outcomes[tx_id] = dataclasses.replace( + service._tx_outcomes[tx_id], timestamp=time.time() - service.TX_OUTCOME_TTL - 1 ) + with pytest.raises(ValueError, match="has not fully terminated"): + service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id=tx_id) + + tx.end_op() + assert service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id=tx_id) == tx_id + + def test_compound_computation_and_release_failures_still_settle_all_sources(self): + """Real cleanup failures can stack: verdict computation may fail while one + source's release hook also raises. Neither failure may skip sibling cleanup or + strand the waiter/ownership marker. + """ + from unittest.mock import Mock, patch + + import nvflare.fuel.f3.streaming.download_service as ds_module - # the live retry is unaffected: no stale terminal outcome surfaces - assert service.get_transaction_outcome("TX-RACE") is None + service = make_isolated_download_service() + service._tx_monitor = Mock() + cell = Mock() + tx_id = service.new_transaction(cell=cell, timeout=10.0, num_receivers=1) + first = _stub_obj() + second = _stub_obj() + first.release = Mock(side_effect=RuntimeError("first release failed")) + service.add_object(tx_id, first) + service.add_object(tx_id, second) + waiter = service.get_transfer_waiter(tx_id) + + with patch.object(ds_module, "compute_transfer_outcome", side_effect=RuntimeError("persistent failure")): + service.delete_transaction(tx_id) + + outcome = waiter.wait(timeout=5.0) + assert outcome is not None and outcome.reason == TransferOutcomeReason.COMPUTATION_FAILED + first.release.assert_called_once_with() + assert second.released, "one raising release must not skip sibling cleanup" + assert first.transaction_done_calls == [(tx_id, TransactionDoneStatus.DELETED)] + assert second.transaction_done_calls == [(tx_id, TransactionDoneStatus.DELETED)] + with service._outcome_lock: + assert tx_id not in service._outcome_owners + with service._tx_lock: + assert tx_id not in service._terminating_txs + + def test_receipt_ttl_starts_at_recording_not_at_verdict(self): + """P1 pin: the outcome timestamp is captured before the settlement callbacks + run -- a settlement slower than TX_OUTCOME_TTL used to record a receipt that + was EXPIRED at birth, which the inline expiry then handed to a same-id + constructor in the gap before the terminator's marker sync. Receipts are now + re-stamped at recording time ("kept 30 min" from recording).""" + from unittest.mock import Mock + + service = make_isolated_download_service() + service._tx_monitor = Mock() + service.TX_OUTCOME_TTL = 0.2 + cell = Mock() + + def slow_done_cb(tid, status, base_objs, **kw): + time.sleep(0.4) # slower than the (patched) TTL + + service.new_transaction( + cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-AGED", transaction_done_cb=slow_done_cb + ) + service.delete_transaction("TX-AGED") + + # the receipt is fresh from RECORDING: queryable, and it excludes the id + assert service.get_transaction_outcome("TX-AGED") is not None + with pytest.raises(ValueError, match="already in use"): + service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-AGED") + + def test_aged_receipt_with_leaked_op_cannot_expose_id(self): + """P1 pin (reviewer's exact sequence): slow settlement past the TTL WITH a + drain-leaked operation -- the id must stay excluded end to end; before the + marker moved to _delete_tx, a constructor could slip between ownership + consumption and the marker sync, and the old marker then overlapped the live + replacement.""" + from unittest.mock import Mock + + import nvflare.fuel.f3.streaming.download_service as ds_module + + service = make_isolated_download_service() + service._tx_monitor = Mock() + service.TX_OUTCOME_TTL = 0.2 + cell = Mock() + service.new_transaction( + cell=cell, + timeout=10.0, + num_receivers=1, + tx_id="TX-AGEDLEAK", + transaction_done_cb=lambda tid, status, base_objs, **kw: time.sleep(0.4), + ) + with service._tx_lock: + tx = service._tx_table["TX-AGEDLEAK"] + assert tx.begin_op() # the leak + + saved = ds_module.OP_DRAIN_TIMEOUT + ds_module.OP_DRAIN_TIMEOUT = 0.1 + try: + service.delete_transaction("TX-AGEDLEAK") finally: - service.shutdown() + ds_module.OP_DRAIN_TIMEOUT = saved + + # settlement complete, receipt aged past the (patched) TTL -- the leaked op + # still excludes the id, and no overlap state can form + time.sleep(0.25) + with pytest.raises(ValueError, match="has not fully terminated"): + service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-AGEDLEAK") + with service._tx_lock: + assert service._terminating_txs.get("TX-AGEDLEAK") is tx + + tx.end_op() + assert service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-AGEDLEAK") == "TX-AGEDLEAK" + + def test_marker_installed_at_unlink_for_every_terminator(self): + """The termination marker is installed by _delete_tx itself, so EVERY + terminator's id is covered for its whole settlement window -- not just + shutdown's.""" + from unittest.mock import Mock + + service = make_isolated_download_service() + service._tx_monitor = Mock() + cell = Mock() + cb_entered = threading.Event() + cb_release = threading.Event() + + def gated_done_cb(tid, status, base_objs, **kw): + cb_entered.set() + cb_release.wait(10.0) + + service.new_transaction( + cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-UNLINK", transaction_done_cb=gated_done_cb + ) + deleter = threading.Thread(target=service.delete_transaction, args=("TX-UNLINK",)) + try: + deleter.start() + assert cb_entered.wait(5.0) + with service._tx_lock: + assert "TX-UNLINK" in service._terminating_txs, "marker must exist mid-settlement on the delete path" + with pytest.raises(ValueError, match="has not fully terminated"): + service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-UNLINK") + finally: + cb_release.set() + deleter.join(5.0) + assert not deleter.is_alive() + + def test_shutdown_window_rejects_id_during_clean_settlement(self): + """P1 pin: the shutdown pre-registration used to be releasable whenever + _active_ops == 0 -- but settlement callbacks are not operations, so a same-id + constructor could pop the marker MID-transaction_done and register while the + old outcome_cb was still about to fire under the id. The marker now needs + settlement_complete AND zero operations to release.""" + from unittest.mock import Mock + + service = make_isolated_download_service() + service._tx_monitor = Mock() + cell = Mock() + cb_entered = threading.Event() + cb_release = threading.Event() + + def gated_done_cb(tid, status, base_objs, **kw): + cb_entered.set() + cb_release.wait(10.0) + + service.new_transaction( + cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-CLEANSET", transaction_done_cb=gated_done_cb + ) + # NOTE: no in-flight operation -- the clean-settlement case the ops-only + # release condition missed + shutter = threading.Thread(target=service.shutdown) + try: + shutter.start() + assert cb_entered.wait(5.0) # settlement mid-callbacks, _active_ops == 0 + + with pytest.raises(ValueError, match="has not fully terminated"): + service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-CLEANSET") + finally: + cb_release.set() + shutter.join(5.0) + assert not shutter.is_alive() + + # settlement complete, no operations: the id frees + assert service.new_transaction(cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-CLEANSET") == "TX-CLEANSET" + + def test_waiter_during_shutdown_settlement_window_resolves_immediately(self): + """shutdown() clears ownership atomically with the table teardown, so a + get_transfer_waiter call landing in the deferred-settlement window finds the + id unknown and resolves None IMMEDIATELY -- it never parks, so it can never + be stranded by the settlements still running.""" + from unittest.mock import Mock + + service = make_isolated_download_service() + service._tx_monitor = Mock() + cell = Mock() + cb_entered = threading.Event() + cb_release = threading.Event() + + def gated_done_cb(tid, status, base_objs, **kw): + cb_entered.set() + cb_release.wait(10.0) + + service.new_transaction( + cell=cell, timeout=10.0, num_receivers=1, tx_id="TX-WIN", transaction_done_cb=gated_done_cb + ) + shutter = threading.Thread(target=service.shutdown) + try: + shutter.start() + assert cb_entered.wait(5.0) # locked teardown done; deferred settlement in flight + + waiter = service.get_transfer_waiter("TX-WIN") + assert waiter.done() and waiter.outcome is None, "window waiter must resolve immediately" + finally: + cb_release.set() + shutter.join(5.0) + assert not shutter.is_alive() + + def test_inflight_budget_pass_drains_before_settlement(self): + """The monitor re-checks table identity, releases _tx_lock, then runs budget + enforcement outside it. Settlement drains such in-flight operations (activity + gate) before snapshotting the outcome: the enforcement's verdicts are COUNTED + in the receipt, and none of its emissions land after the transaction settled.""" + from unittest.mock import Mock + + service = make_isolated_download_service() + service._tx_monitor = Mock() + cell = Mock() + events = [] + cb_entered = threading.Event() + cb_release = threading.Event() + + def gated_downloaded_to_one(receiver, status): + cb_entered.set() + cb_release.wait(10.0) + events.append("old_downloaded_to_one") + + service.new_transaction( + cell=cell, + timeout=1000.0, + num_receivers=1, + tx_id="TX-BUDGET", + receiver_ids=("r1",), + receiver_acquire_timeout=5.0, + progress_cb=lambda **kw: events.append(("old_progress", kw.get("state"))), + outcome_cb=lambda outcome: events.append("outcome_recorded"), + ) + obj = _stub_obj() + obj.downloaded_to_one = gated_downloaded_to_one + service.add_object("TX-BUDGET", obj) + + monitor = threading.Thread(target=run_monitor_once, args=(service,), kwargs={"now": time.time() + 100.0}) + deleter = threading.Thread(target=service.delete_transaction, args=("TX-BUDGET",)) + try: + monitor.start() + assert cb_entered.wait(5.0), "budget enforcement must have finalized r1 as FAILED" + + # settlement lands while enforcement is mid-flight: it must drain first + deleter.start() + deleter.join(0.5) + assert deleter.is_alive(), "settlement must not proceed while a budget pass is in flight" + finally: + cb_release.set() + monitor.join(5.0) + deleter.join(5.0) + assert not deleter.is_alive() + + # the enforcement's emissions all precede outcome recording, and its verdict + # was counted: r1 books as FAILED in the receipt + recorded_at = events.index("outcome_recorded") + assert events.index("old_downloaded_to_one") < recorded_at + assert events.index(("old_progress", TransferProgressState.FAILED)) < recorded_at + outcome = service.get_transaction_outcome("TX-BUDGET") + assert outcome is not None and not outcome.completed + assert outcome.refs[0].receiver_statuses.get("r1") == DownloadStatus.FAILED def test_unknown_transaction_has_no_outcome(self): service = make_isolated_download_service() @@ -366,7 +1009,11 @@ def test_shutdown_clears_outcomes_and_stops_recording(self): service.shutdown() assert service.get_transaction_outcome(tx.tid) is None - # a monitor iteration that was mid-termination during shutdown cannot repopulate + # a monitor iteration that was mid-termination during shutdown cannot repopulate: + # shutdown cleared outcome ownership, so the late recorder's tx no longer owns + # the slot and its outcome drops + from unittest.mock import Mock + late = compute_transfer_outcome("T-LATE", TransactionDoneStatus.FINISHED, 1, [], time.time()) - service._record_outcome(late) + service._record_outcome(late, tx=Mock()) assert service.get_transaction_outcome("T-LATE") is None diff --git a/tests/unit_test/fuel/f3/streaming/transfer_waiter_test.py b/tests/unit_test/fuel/f3/streaming/transfer_waiter_test.py new file mode 100644 index 0000000000..a9aed72948 --- /dev/null +++ b/tests/unit_test/fuel/f3/streaming/transfer_waiter_test.py @@ -0,0 +1,232 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. 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. + +"""Tests for the awaitable transfer facade (TransferWaiter). + +TransferWaiter is the "returns == delivered" primitive the upper layers (executor backends, +trainer engine) consume: wait() blocks -- event-driven, resolved inside the outcome-recording +path -- until the aggregate TransferOutcome is recorded; COMPLETED only when every expected +receiver succeeded. It composes with (never replaces) transaction_done_cb / outcome_cb. +""" + +import threading +import time +from unittest.mock import Mock, patch + +from nvflare.fuel.f3.streaming import download_service as ds_module +from nvflare.fuel.f3.streaming.download_service import DownloadStatus, TransferWaiter +from tests.unit_test.fuel.f3.streaming.download_test_utils import MockDownloadable, confirm_request +from tests.unit_test.fuel.f3.streaming.download_test_utils import make_service_no_monitor as _make_service +from tests.unit_test.fuel.f3.streaming.download_test_utils import pull_to_terminal as _pull_to_terminal +from tests.unit_test.fuel.f3.streaming.download_test_utils import run_monitor_once, serve_nonce + + +def _new_tx(service, num_receivers=1, timeout=10.0): + tx_id = service.new_transaction(cell=Mock(), timeout=timeout, num_receivers=num_receivers) + rid = service.add_object(tx_id, MockDownloadable([b"chunk"])) + return tx_id, rid + + +class TestTransferWaiter: + def test_wait_blocks_until_outcome_and_returns_delivered(self): + # the load-bearing property: a thread blocked in wait() is released by outcome + # recording itself (event-driven, no polling), and gets the COMPLETED certificate + service = _make_service() + tx_id, rid = _new_tx(service) + waiter = service.get_transfer_waiter(tx_id) + assert not waiter.done() + + results = [] + t = threading.Thread(target=lambda: results.append(waiter.wait(timeout=10.0))) + t.start() + + terminal = _pull_to_terminal(service, rid, "r1", confirm_capable=True) + service._handle_download(confirm_request(rid, "r1", DownloadStatus.SUCCESS, serve_nonce(terminal))) + run_monitor_once(service, now=time.time()) + + t.join(5.0) + assert not t.is_alive() + assert len(results) == 1 and results[0] is not None + assert results[0].completed + assert results[0].tx_id == tx_id + + def test_waiter_after_termination_resolves_immediately(self): + service = _make_service() + tx_id, rid = _new_tx(service) + _pull_to_terminal(service, rid, "r1") + run_monitor_once(service, now=time.time()) + + waiter = service.get_transfer_waiter(tx_id) + assert waiter.done() + outcome = waiter.wait(timeout=0) + assert outcome is not None and outcome.completed + + def test_wait_timeout_returns_none_then_resolves(self): + service = _make_service() + tx_id, rid = _new_tx(service) + waiter = service.get_transfer_waiter(tx_id) + + assert waiter.wait(timeout=0.05) is None + assert waiter.outcome is None + + _pull_to_terminal(service, rid, "r1") + run_monitor_once(service, now=time.time()) + outcome = waiter.wait(timeout=5.0) + assert outcome is not None and outcome.completed + + def test_failed_outcome_is_returned_not_masked(self): + service = _make_service() + tx_id, rid = _new_tx(service) + waiter = service.get_transfer_waiter(tx_id) + + terminal = _pull_to_terminal(service, rid, "r1", confirm_capable=True) + service._handle_download( + confirm_request(rid, "r1", DownloadStatus.FAILED, serve_nonce(terminal)) + ) # receiver truth: finalization failed + run_monitor_once(service, now=time.time()) + + outcome = waiter.wait(timeout=5.0) + assert outcome is not None + assert not outcome.completed + + def test_linger_applies_to_finished_outcomes_only(self): + # linger preserves the process/tombstone window so lost terminal replies can be + # replayed: it applies to ANY FINISHED outcome (completed or receiver-failed -- + # the served-SUCCESS receivers of a partial fan-out are exactly who needs it), + # and not to timed-out transactions (no receiver is retrying a served reply). + service = _make_service() + + # completed FINISHED: linger applied + tx_id, rid = _new_tx(service) + _pull_to_terminal(service, rid, "r1") + run_monitor_once(service, now=time.time()) + waiter = service.get_transfer_waiter(tx_id) + with patch.object(ds_module.time, "sleep") as mock_sleep: + outcome = waiter.wait(timeout=5.0, linger=0.2) + assert outcome.completed + mock_sleep.assert_called_once_with(0.2) + + # receiver-failed FINISHED: linger still applied (partial fan-out healing window) + tx_id2, rid2 = _new_tx(service) + terminal2 = _pull_to_terminal(service, rid2, "r1", confirm_capable=True) + service._handle_download(confirm_request(rid2, "r1", DownloadStatus.FAILED, serve_nonce(terminal2))) + run_monitor_once(service, now=time.time()) + waiter2 = service.get_transfer_waiter(tx_id2) + with patch.object(ds_module.time, "sleep") as mock_sleep: + outcome2 = waiter2.wait(timeout=5.0, linger=0.2) + assert not outcome2.completed + mock_sleep.assert_called_once_with(0.2) + + # TIMEOUT termination: no linger + tx_id3, rid3 = _new_tx(service) + service._tx_table[tx_id3].last_active_time = time.time() - 100.0 + run_monitor_once(service, now=time.time()) + waiter3 = service.get_transfer_waiter(tx_id3) + with patch.object(ds_module.time, "sleep") as mock_sleep: + outcome3 = waiter3.wait(timeout=5.0, linger=5.0) + assert outcome3 is not None and not outcome3.completed + mock_sleep.assert_not_called() + + def test_waiter_for_unknown_tx_resolves_immediately(self): + # invariant: waiters can never hang -- nothing will ever record an outcome for an + # unknown (or already-forgotten) transaction id + service = _make_service() + waiter = service.get_transfer_waiter("no-such-tx") + assert waiter.done() + assert waiter.wait(timeout=0) is None + assert "no-such-tx" not in service._tx_waiters + + def test_waiter_after_shutdown_resolves_immediately(self): + service = _make_service() + tx_id, _ = _new_tx(service) + service.shutdown() + + waiter = service.get_transfer_waiter(tx_id) + assert waiter.done() + assert waiter.wait(timeout=0) is None + assert not service._tx_waiters + + def test_shutdown_unblocks_waiters_with_none(self): + service = _make_service() + tx_id, _ = _new_tx(service) + waiter = service.get_transfer_waiter(tx_id) + + results = [] + t = threading.Thread(target=lambda: results.append(waiter.wait(timeout=10.0))) + t.start() + service.shutdown() + t.join(5.0) + + assert not t.is_alive(), "shutdown must never leave a waiter hanging" + assert results == [None] + + def test_acquired_receivers_reflects_first_pulls(self): + service = _make_service() + tx_id, rid = _new_tx(service, num_receivers=2) + assert service.get_acquired_receivers(tx_id) == set() + + _pull_to_terminal(service, rid, "r1") + assert service.get_acquired_receivers(tx_id) == {"r1"} + + def test_budget_failure_resolves_waiter_in_bounded_time(self): + # end-to-end with the budgets: a stalled receiver's budget failure terminates the tx and + # releases the waiter -- the producer is never pinned to the full TTL + service = _make_service() + tx_id = service.new_transaction(cell=Mock(), timeout=1000.0, num_receivers=2, receiver_idle_timeout=5.0) + rid = service.add_object(tx_id, MockDownloadable([b"chunk"])) + waiter = service.get_transfer_waiter(tx_id) + + _pull_to_terminal(service, rid, "healthy") + _pull_to_terminal(service, rid, "stalled", confirm_capable=True) # confirm never arrives + run_monitor_once(service, now=time.time() + 30.0) + + outcome = waiter.wait(timeout=5.0) + assert outcome is not None + assert not outcome.completed + assert outcome.refs[0].receiver_statuses["stalled"] == DownloadStatus.FAILED + + def test_wait_returns_only_after_callbacks_and_release(self): + # P1 pin: an upper layer that stops the producer when wait() returns must not be + # able to preempt the callback chain or source release + service = _make_service() + settled = [] + tx_id = service.new_transaction( + cell=Mock(), + timeout=10.0, + num_receivers=1, + transaction_done_cb=lambda *a, **kw: settled.append("done_cb"), + outcome_cb=lambda outcome: settled.append("outcome_cb"), + ) + obj = MockDownloadable([b"chunk"]) + rid = service.add_object(tx_id, obj) + waiter = service.get_transfer_waiter(tx_id) + + results = [] + t = threading.Thread(target=lambda: results.append(waiter.wait(timeout=10.0))) + t.start() + _pull_to_terminal(service, rid, "r1") + run_monitor_once(service, now=time.time()) + t.join(5.0) + + assert results and results[0] is not None and results[0].completed + # by the time wait() returned, the transaction was fully settled + assert settled == ["done_cb", "outcome_cb"] + assert obj.released + + def test_waiter_is_a_transfer_waiter(self): + service = _make_service() + tx_id, _ = _new_tx(service) + waiter = service.get_transfer_waiter(tx_id) + assert isinstance(waiter, TransferWaiter) + assert waiter.transaction_id == tx_id