diff --git a/desktop/scripts/verify-service-contracts.ts b/desktop/scripts/verify-service-contracts.ts index aceefee68..e0b237059 100644 --- a/desktop/scripts/verify-service-contracts.ts +++ b/desktop/scripts/verify-service-contracts.ts @@ -137,8 +137,11 @@ const NOTIFY_LITERAL_RE = /(?:(?:notify|emit|publish|broadcast)\(\s*|notifyWithShutdown\(\s*ctx\s*,\s*)"([^"]+)"/gi const NOTIFY_PREFIX_CONCAT_RE = /(?:(?:notify|emit|publish|broadcast)\(\s*|notifyWithShutdown\(\s*ctx\s*,\s*)"([^"]+)"\s*\+/i +// Broadcast(ctx, frame) is HTTP fan-out, not a JSON-RPC notification. Keep +// literal Broadcast methods above, but do not treat an identifier argument to +// Broadcast as an unresolved JSON-RPC method. const NOTIFY_IDENT_RE = - /(?:(?:notify|emit|publish|broadcast)\(\s*|notifyWithShutdown\(\s*ctx\s*,\s*)([A-Za-z_]\w*)\s*,/gi + /(?:(?:notify|emit|publish)\(\s*|notifyWithShutdown\(\s*ctx\s*,\s*)([A-Za-z_]\w*)\s*,/gi /** A `| `nvpair-ui-broker` | 0.37.0 | … |` row — forbidden in hand-maintained docs. */ const DOC_COMPONENT_ROW_RE = /^\|\s*`([a-z0-9-]+)`\s*\|\s*(\d+\.\d+\.\d+)\s*\|/gm diff --git a/desktop/src/electron/service-bridge/modular-state.ts b/desktop/src/electron/service-bridge/modular-state.ts index c7bbb8f9e..0786ad174 100644 --- a/desktop/src/electron/service-bridge/modular-state.ts +++ b/desktop/src/electron/service-bridge/modular-state.ts @@ -38,6 +38,7 @@ import type { Workload, WorkloadLiveEngine, WorkloadLiveProgress, + WorkloadRemoval, WorkloadState } from '@/shared/types/workloads' import type { WorkerState, WorkerStatus } from '@/shared/types/workers' @@ -55,7 +56,7 @@ import { parseEngineCapabilities } from '@/shared/utils/engines' import { engineProgressKey } from '@/shared/utils/engine-progress' -import { workloadKey } from '@/shared/utils/workloads' +import { workloadKey, workloadMatchesRemoval } from '@/shared/utils/workloads' import { currentPlatform, platformDisplayName } from '@/shared/utils/platform' import { emitBridgePush } from './broadcaster' import { createStructuredLogger } from '@/shared/utils/log' @@ -673,6 +674,7 @@ function parseWorkload(value: JsonValue | undefined): Workload | null { engine, state: stateValue, originatedFrom: nullableStringValue(obj.originatedFrom), + runId: nullableStringValue(obj.runId), createdAt: numberValue(obj.createdAt), startedAt: nullableNumberValue(obj.startedAt), completedAt: nullableNumberValue(obj.completedAt), @@ -1520,6 +1522,13 @@ class ModularBridgeState { private workloads = new Map() // Exists only during one epoch snapshot fetch, including removal tombstones. private workloadRecovery: Set | null = null + /** + * Broker removals seen while a recovery snapshot is fetched. A removal + * names a job by `(originatedFrom, id)` only, so it is kept as itself + * rather than as keys: it must also retire a matching job that is not in + * the catalog yet but is in the older snapshot. + */ + private workloadRecoveryRemovals: WorkloadRemoval[] = [] private logs: LogEntry[] = [] // Per-engine bound proxy port reported by the broker. 0 = not reported yet; // we never fabricate a default — an unknown port surfaces as null, not a @@ -1948,7 +1957,7 @@ class ModularBridgeState { */ seedWorkloads(workloads: Workload[]): WsInvokeResponse<'workloads:get-initial'> { for (const workload of workloads) { - const key = workloadKey(workload.originatedFrom, workload.id) + const key = workloadKey(workload) if (!this.workloads.has(key)) this.workloads.set(key, workload) } return this.getWorkloads() @@ -1958,6 +1967,7 @@ class ModularBridgeState { beginWorkloadRecovery(): Set { const snapshot = new Set() this.workloadRecovery = snapshot + this.workloadRecoveryRemovals = [] return snapshot } @@ -1965,18 +1975,19 @@ class ModularBridgeState { endWorkloadRecovery(snapshot: Set, workloads?: Workload[]): void { if (this.workloadRecovery !== snapshot) return this.workloadRecovery = null + const removals = this.workloadRecoveryRemovals + this.workloadRecoveryRemovals = [] if (!workloads) { snapshot.clear() return } - const baseline = new Map( - workloads.map(workload => [workloadKey(workload.originatedFrom, workload.id), workload]) - ) + const baseline = new Map(workloads.map(workload => [workloadKey(workload), workload])) for (const [key, workload] of this.workloads) { if (!snapshot.has(key) && !baseline.has(key)) this.evictWorkload(key, workload) } for (const [key, workload] of baseline) { if (snapshot.has(key)) continue + if (removals.some(removal => workloadMatchesRemoval(workload, removal))) continue this.workloads.set(key, workload) emitBridgePush('workloads:upsert', workload) } @@ -1988,22 +1999,29 @@ class ModularBridgeState { const obj = objectValue(params) const workload = parseWorkload(obj?.workloadInfo) if (!workload) return - const key = workloadKey(workload.originatedFrom, workload.id) + const key = workloadKey(workload) this.workloadRecovery?.add(key) this.workloads.set(key, workload) emitBridgePush('workloads:upsert', workload) } - /** Relay a broker `workloads:remove` (`{ workloadId, originatedFrom? }`) into the catalog + UI. */ + /** + * Relay a broker `workloads:remove` (`{ workloadId, originatedFrom? }`) into + * the catalog + UI. It names no engine and no run, so, like the broker's own + * store, it drops every job matching the pair. + */ removeWorkloadFromParams(params: JsonValue | undefined): void { const obj = objectValue(params) const workloadId = stringValue(obj?.workloadId) if (!workloadId) return - const originatedFrom = nullableStringValue(obj?.originatedFrom) - const key = workloadKey(originatedFrom, workloadId) - this.workloadRecovery?.add(key) - this.workloads.delete(key) - emitBridgePush('workloads:remove', { workloadId, originatedFrom }) + const removal = { workloadId, originatedFrom: nullableStringValue(obj?.originatedFrom) } + if (this.workloadRecovery) this.workloadRecoveryRemovals.push(removal) + for (const [key, workload] of this.workloads) { + if (!workloadMatchesRemoval(workload, removal)) continue + this.workloadRecovery?.add(key) + this.workloads.delete(key) + } + emitBridgePush('workloads:remove', removal) } /** @@ -2027,7 +2045,9 @@ class ModularBridgeState { this.workloads.delete(key) emitBridgePush('workloads:remove', { workloadId: workload.id, - originatedFrom: workload.originatedFrom + originatedFrom: workload.originatedFrom, + engine: workload.engine, + runId: workload.runId ?? null }) } diff --git a/desktop/src/electron/tray-badge.ts b/desktop/src/electron/tray-badge.ts index 1f184bebd..078a99481 100644 --- a/desktop/src/electron/tray-badge.ts +++ b/desktop/src/electron/tray-badge.ts @@ -26,7 +26,7 @@ import { type TrayBadgeInk, type TrayBadgeMap } from '@/shared/utils/tray-badge-map' -import { workloadKey } from '@/shared/utils/workloads' +import { workloadKey, workloadMatchesRemoval } from '@/shared/utils/workloads' export interface TrayBadgeSink { setIcon: (map: TrayBadgeMap) => void @@ -66,7 +66,7 @@ export function startTrayBadge( ): () => void { const catalog = new Map() for (const workload of initial) { - catalog.set(workloadKey(workload.originatedFrom, workload.id), workload) + catalog.set(workloadKey(workload), workload) } let signature = '' @@ -105,9 +105,11 @@ export function startTrayBadge( const unsubscribe = subscribePush(event => { if (event.channel === 'workloads:upsert') { - catalog.set(workloadKey(event.payload.originatedFrom, event.payload.id), event.payload) + catalog.set(workloadKey(event.payload), event.payload) } else if (event.channel === 'workloads:remove') { - catalog.delete(workloadKey(event.payload.originatedFrom, event.payload.workloadId)) + for (const [key, workload] of catalog) { + if (workloadMatchesRemoval(workload, event.payload)) catalog.delete(key) + } } else { return } diff --git a/desktop/src/shared/types/workloads.ts b/desktop/src/shared/types/workloads.ts index f969fafba..15fb3156e 100644 --- a/desktop/src/shared/types/workloads.ts +++ b/desktop/src/shared/types/workloads.ts @@ -13,11 +13,17 @@ export interface Workload { state: WorkloadState /** * Owner/origin node of the workload — the node whose proxy received the - * request. This is the identity half of the backend's `(originatedFrom, id)` - * global catalog key (workload ids are a per-node proxy counter, so they - * collide across nodes; `originatedFrom` disambiguates). + * request. Part of the backend's `(originatedFrom, engine, runId, id)` + * catalog key: workload ids are a per-proxy counter, so they collide across + * nodes and engines. */ originatedFrom: string | null + /** + * The proxy run that numbered this job. A restarted proxy counts from 1 + * again; the run id tells its jobs from the earlier run's. Absent from a + * producer that predates it. + */ + runId?: string | null /** * Node the workload was routed to / scheduled on — where it actually ran, * as opposed to `originatedFrom` (where it came from). Optional/additive: @@ -131,3 +137,11 @@ export interface WorkloadLiveProgress { */ engineProgress?: WorkloadLiveEngine | null } + +/** A `workloads:remove` push. The broker's carries no engine and no run id. */ +export interface WorkloadRemoval { + workloadId: string + originatedFrom: string | null + engine?: EngineType | null + runId?: string | null +} diff --git a/desktop/src/shared/types/ws-channels.ts b/desktop/src/shared/types/ws-channels.ts index 0a96a7814..7e66c85f2 100644 --- a/desktop/src/shared/types/ws-channels.ts +++ b/desktop/src/shared/types/ws-channels.ts @@ -39,7 +39,7 @@ import type { NodeItemMetrics, NodeInferenceMetricsSnapshot } from '@/shared/types/metrics' -import type { Workload, WorkloadLiveProgress } from '@/shared/types/workloads' +import type { Workload, WorkloadLiveProgress, WorkloadRemoval } from '@/shared/types/workloads' import type { WorkerStatus } from '@/shared/types/workers' import type { AvailableNode, @@ -171,7 +171,7 @@ export interface WsPushChannelMap { // workload ids are a per-node proxy counter (the catalog is keyed by the // (originatedFrom, id) pair). 'workloads:upsert': Workload - 'workloads:remove': { workloadId: string; originatedFrom: string | null } + 'workloads:remove': WorkloadRemoval /** * The broker's live view of a running job: what PAIR measured of the * client's stream (elapsed time, chunks, bytes, time to first byte) plus diff --git a/desktop/src/shared/utils/workloads.ts b/desktop/src/shared/utils/workloads.ts index 34eaae491..46f303ff9 100644 --- a/desktop/src/shared/utils/workloads.ts +++ b/desktop/src/shared/utils/workloads.ts @@ -1,24 +1,53 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -import type { Workload, WorkloadLiveProgress } from '@/shared/types/workloads' +import type { Workload, WorkloadLiveProgress, WorkloadRemoval } from '@/shared/types/workloads' import type { EngineType } from '@/shared/types/engines' +/** The fields that identify one job. */ +interface WorkloadIdentity { + originatedFrom: string | null + engine: EngineType | null + runId?: string | null + id: string +} + /** - * Stable catalog key for a workload. + * Stable catalog key for a workload: the backend's own identity, + * `(originatedFrom, engine, runId, id)`. * - * The backend's catalog is keyed by `(originatedFrom, engine, runId, id)`, but - * the `workloads:remove` push carries only `(workloadId, originatedFrom)` — and - * the broker's own `Store.Remove` drops every record matching that pair — so - * `(originatedFrom, id)` is the only key a subscribe client can maintain - * consistently across upsert and remove. Each node's proxy assigns workload ids - * from its own monotonic counter, so ids collide across nodes; `originatedFrom` - * (the origin node) disambiguates. Mirror that here so a remote node's job never - * overwrites a local one that happens to share an id. The `\u0000` separator - * cannot appear in a host id or proxy counter, so the key is unambiguous. + * Every proxy numbers its jobs from 1, and numbers again from 1 when it + * restarts. The origin tells nodes apart, the engine tells an Ollama job "3" + * from an oMLX job "3" on one node, and the run id tells a job from the one a + * restarted proxy gave the same number. Keyed on fewer fields, one job's card + * would overwrite another's. The `\u0000` separator cannot appear in any of + * the fields, so the key is unambiguous. + */ +export function workloadKey(workload: WorkloadIdentity): string { + return [ + workload.originatedFrom ?? '', + workload.engine ?? '', + workload.runId ?? '', + workload.id + ].join('\u0000') +} + +/** + * Whether a removal names this workload. A removal from the broker carries + * only `(workloadId, originatedFrom)` and, like the broker's own store, drops + * every job matching that pair. One raised by the desktop itself names the + * engine and run too — `runId: null` meaning "the job with no run", not "any + * run" — so it drops exactly one job. Only an absent field is a wildcard. */ -export function workloadKey(originatedFrom: string | null, id: string): string { - return `${originatedFrom ?? ''}\u0000${id}` +export function workloadMatchesRemoval( + workload: WorkloadIdentity, + removal: WorkloadRemoval +): boolean { + if (workload.id !== removal.workloadId) return false + if ((workload.originatedFrom ?? '') !== (removal.originatedFrom ?? '')) return false + if (removal.engine !== undefined && workload.engine !== removal.engine) return false + if (removal.runId !== undefined && (workload.runId ?? null) !== removal.runId) return false + return true } /** diff --git a/desktop/src/ui/components/Workloads/WorkloadItemCard.tsx b/desktop/src/ui/components/Workloads/WorkloadItemCard.tsx index 2421e2788..9436f1565 100644 --- a/desktop/src/ui/components/Workloads/WorkloadItemCard.tsx +++ b/desktop/src/ui/components/Workloads/WorkloadItemCard.tsx @@ -4,7 +4,7 @@ import { useMemo, memo } from 'react' import { Card, Flex, Stack, Text } from '@nvidia/foundations-react-core' import type { Workload } from '@/shared/types/workloads' -import { workloadExecutionNodeId } from '@/shared/utils/workloads' +import { workloadExecutionNodeId, workloadKey } from '@/shared/utils/workloads' import { useNodesStore } from '@/ui/stores/nodes.store' import { useWorkloadLiveProgress } from '@/ui/stores/workloads.store' import { formatModelDisplayName } from '@/ui/utils/format-model-display-name' @@ -148,8 +148,7 @@ function WorkloadItemCard({ workload }: { workload: Workload }) { className={className} density="compact" style={{ direction: 'ltr' }} - data-workload-id={workload.id} - data-workload-origin={workload.originatedFrom ?? ''} + data-workload-key={workloadKey(workload)} attributes={{ CardContent: { className: 'workload-card-content' } }} >
diff --git a/desktop/src/ui/components/Workloads/WorkloadListView.tsx b/desktop/src/ui/components/Workloads/WorkloadListView.tsx index bfcaa28cf..064cc45ef 100644 --- a/desktop/src/ui/components/Workloads/WorkloadListView.tsx +++ b/desktop/src/ui/components/Workloads/WorkloadListView.tsx @@ -63,9 +63,7 @@ function WorkloadListView({ filter, setFilter }: WorkloadListViewProps) { const filterKey = Object.entries(filter) .map(([_, value]) => `${value}`) .join('-') - const workloadsKey = currentWorkloads - .map(w => workloadKey(w.originatedFrom, w.id)) - .join('-') + const workloadsKey = currentWorkloads.map(w => workloadKey(w)).join('-') return `${filterKey}-${workloadsKey}` }, [filter, currentWorkloads]) @@ -110,7 +108,7 @@ function WorkloadListView({ filter, setFilter }: WorkloadListViewProps) { data-workload-list-content > {currentWorkloads.map(workload => { - const key = workloadKey(workload.originatedFrom, workload.id) + const key = workloadKey(workload) return (
diff --git a/desktop/src/ui/components/Workloads/WorkloadNodeConnections.tsx b/desktop/src/ui/components/Workloads/WorkloadNodeConnections.tsx index 60d88a878..dac25d9c3 100644 --- a/desktop/src/ui/components/Workloads/WorkloadNodeConnections.tsx +++ b/desktop/src/ui/components/Workloads/WorkloadNodeConnections.tsx @@ -50,14 +50,6 @@ function connectionsEqual(a: readonly Connection[], b: readonly Connection[]): b return true } -// Build a key that uniquely identifies a workload's DOM anchor. Workload ids are -// a per-node proxy counter and collide across nodes, so we key on the -// (origin, id) pair the same way `workloadKey` does. The null byte can't appear -// in either value, so it is a safe separator. -function anchorKey(origin: string, id: string): string { - return `${origin}\u0000${id}` -} - function computeConnections(svg: SVGSVGElement, workloads: readonly Workload[]): Connection[] { const svgRect = svg.getBoundingClientRect() const result: Connection[] = [] @@ -66,10 +58,9 @@ function computeConnections(svg: SVGSVGElement, workloads: readonly Workload[]): // frame. With N running jobs the old path issued 2 * N selector queries every // animation frame; this issues two querySelectorAll calls total. const workloadElByKey = new Map() - for (const el of document.querySelectorAll('[data-workload-id]')) { - const id = el.getAttribute('data-workload-id') - if (id === null) continue - workloadElByKey.set(anchorKey(el.getAttribute('data-workload-origin') ?? '', id), el) + for (const el of document.querySelectorAll('[data-workload-key]')) { + const key = el.getAttribute('data-workload-key') + if (key !== null) workloadElByKey.set(key, el) } const nodeElById = new Map() for (const el of document.querySelectorAll('[data-node-id]')) { @@ -81,9 +72,7 @@ function computeConnections(svg: SVGSVGElement, workloads: readonly Workload[]): const executionNodeId = workloadExecutionNodeId(workload) if (!executionNodeId) continue - const workloadEl = workloadElByKey.get( - anchorKey(workload.originatedFrom ?? '', workload.id) - ) + const workloadEl = workloadElByKey.get(workloadKey(workload)) const nodeEl = nodeElById.get(executionNodeId) if (!workloadEl || !nodeEl) continue @@ -126,7 +115,7 @@ function computeConnections(svg: SVGSVGElement, workloads: readonly Workload[]): const color = getWorkloadStateColor(workload.state) result.push({ - workloadId: workloadKey(workload.originatedFrom, workload.id), + workloadId: workloadKey(workload), nodeId: executionNodeId, color: `${WORKLOAD_COLOR_MAP[color]}${workload.state === 'running' ? 'FF' : '88'}`, path, diff --git a/desktop/src/ui/stores/workloads.store.ts b/desktop/src/ui/stores/workloads.store.ts index a23aed885..00b8bbf15 100644 --- a/desktop/src/ui/stores/workloads.store.ts +++ b/desktop/src/ui/stores/workloads.store.ts @@ -3,14 +3,15 @@ import { create } from 'zustand' import { useShallow } from 'zustand/react/shallow' -import type { Workload, WorkloadLiveProgress } from '@/shared/types/workloads' +import type { Workload, WorkloadLiveProgress, WorkloadRemoval } from '@/shared/types/workloads' import deepEqual from '@/ui/utils/deep-equal' import { liveProgressKey, liveProgressKeyIsJob, mergeWorkloadLiveProgress, workloadExecutionNodeId, - workloadKey + workloadKey, + workloadMatchesRemoval } from '@/shared/utils/workloads' import { stateOrder, MAX_HISTORY_ITEMS } from '@/ui/constants/app' @@ -48,7 +49,7 @@ let unsubs: Array<() => void> = [] // buffer pushes and flush them in a single commit on the next animation frame. type PendingOp = | { kind: 'upsert'; workload: Workload } - | { kind: 'remove'; key: string; originatedFrom: string | null; id: string } + | { kind: 'remove'; removal: WorkloadRemoval } | { kind: 'live'; key: string; frame: WorkloadLiveProgress } let pendingOps: PendingOp[] = [] let flushHandle = 0 @@ -76,16 +77,17 @@ export const useWorkloadsStore = create((set, get) => ({ // A job that has already finished has no live line: the // broker retires its live view, so a frame that arrives // after the terminal event must not bring it back. - const workload = (updated ?? prev).get( - workloadKey(op.frame.originatedFrom, op.frame.jobId) + const finished = [...(updated ?? prev).values()].some( + workload => + workload.id === op.frame.jobId && + workload.originatedFrom === op.frame.originatedFrom && + workload.engine === op.frame.engine && + (!op.frame.runId || + !workload.runId || + workload.runId === op.frame.runId) && + (workload.state === 'completed' || workload.state === 'failed') ) - if ( - workload && - workload.engine === op.frame.engine && - (workload.state === 'completed' || workload.state === 'failed') - ) { - continue - } + if (finished) continue const existing = base.get(op.key) const merged = mergeWorkloadLiveProgress(existing, op.frame) if (existing && deepEqual(existing, merged)) continue @@ -104,7 +106,7 @@ export const useWorkloadsStore = create((set, get) => ({ } const base = updated ?? prev if (op.kind === 'upsert') { - const key = workloadKey(op.workload.originatedFrom, op.workload.id) + const key = workloadKey(op.workload) const existing = base.get(key) if (existing && deepEqual(existing, op.workload)) continue if (!updated) updated = new Map(prev) @@ -121,15 +123,35 @@ export const useWorkloadsStore = create((set, get) => ({ ) } } else { - if (!base.has(op.key)) continue + const { removal } = op + const keys = [...base] + .filter(([, w]) => workloadMatchesRemoval(w, removal)) + .map(([k]) => k) + if (keys.length === 0) continue if (!updated) updated = new Map(prev) - updated.delete(op.key) + for (const key of keys) updated.delete(key) if (!updatedLive) updatedLive = new Map(prevLive) - // A removal names no engine, so it retires the live view - // of the job on every engine of that origin. - for (const liveKey of [...updatedLive.keys()]) { - if (liveProgressKeyIsJob(liveKey, op.originatedFrom, op.id)) { - updatedLive.delete(liveKey) + if (removal.engine) { + updatedLive.delete( + liveProgressKey( + removal.originatedFrom, + removal.engine, + removal.workloadId + ) + ) + } else { + // A removal that names no engine retires the live view + // of the job on every engine of that origin. + for (const liveKey of [...updatedLive.keys()]) { + if ( + liveProgressKeyIsJob( + liveKey, + removal.originatedFrom, + removal.workloadId + ) + ) { + updatedLive.delete(liveKey) + } } } } @@ -152,7 +174,7 @@ export const useWorkloadsStore = create((set, get) => ({ // below would be silently undone by rebuilding the map from the baseline // snapshot, which still lists the now-retired job. Record such removes so // the merge can subtract them — these deltas are newer than the snapshot. - const removedDuringInit = new Set() + const removedDuringInit: WorkloadRemoval[] = [] let initializing = true // Subscribe BEFORE fetching the baseline so a `workloads:upsert` that @@ -166,10 +188,9 @@ export const useWorkloadsStore = create((set, get) => ({ pendingOps.push({ kind: 'upsert', workload }) schedule() }), - window.pairApi.workloads.onRemove(({ workloadId, originatedFrom }) => { - const key = workloadKey(originatedFrom, workloadId) - if (initializing) removedDuringInit.add(key) - pendingOps.push({ kind: 'remove', key, originatedFrom, id: workloadId }) + window.pairApi.workloads.onRemove((removal: WorkloadRemoval) => { + if (initializing) removedDuringInit.push(removal) + pendingOps.push({ kind: 'remove', removal }) schedule() }), // The live view of a running job. It is a snapshot, not a @@ -192,7 +213,10 @@ export const useWorkloadsStore = create((set, get) => ({ // Subtract jobs a `workloads:remove` retired during the fetch (that // delta is newer than the snapshot), before overlaying upserts — so a // remove-then-readd of the same key still surfaces the re-add. - for (const key of removedDuringInit) map.delete(key) + for (const [key, workload] of map) { + if (removedDuringInit.some(removal => workloadMatchesRemoval(workload, removal))) + map.delete(key) + } // Overlay upserts that already landed during the fetch (also newer than // the baseline), so seeding never regresses a live transition. for (const [key, workload] of get().workloads) { diff --git a/desktop/tests/modular/epoch-resync.test.ts b/desktop/tests/modular/epoch-resync.test.ts index e2908818a..16ba05636 100644 --- a/desktop/tests/modular/epoch-resync.test.ts +++ b/desktop/tests/modular/epoch-resync.test.ts @@ -361,24 +361,26 @@ describe('broker epoch tracking and resync', () => { }) await flushMicrotasks() const workloads = state.getWorkloads() - expect(workloads['origin\u0000terminal']).toMatchObject({ + expect(workloads['origin\u0000ollama\u0000\u0000terminal']).toMatchObject({ state: 'completed', model: 'snapshot' }) - expect(workloads['origin\u0000absent']).toBeUndefined() - expect(workloads['origin\u0000upsert']).toMatchObject({ + expect(workloads['origin\u0000ollama\u0000\u0000absent']).toBeUndefined() + expect(workloads['origin\u0000ollama\u0000\u0000upsert']).toMatchObject({ state: 'running', model: 'new-push' }) - expect(workloads['origin\u0000removed']).toBeUndefined() + expect(workloads['origin\u0000ollama\u0000\u0000removed']).toBeUndefined() expect(calls.filter(m => m === 'workloads:get-initial')).toHaveLength(1) delete responses['workloads:get-initial'] responses['workloads:get-initial'] = () => Promise.resolve({ workloads: [job('removed', 'completed')] }) noteEpoch('workloads-next') await flushMicrotasks() - expect(state.getWorkloads()['origin\u0000removed']).toMatchObject({ state: 'completed' }) - expect(state.getWorkloads()['origin\u0000upsert']).toBeUndefined() + expect(state.getWorkloads()['origin\u0000ollama\u0000\u0000removed']).toMatchObject({ + state: 'completed' + }) + expect(state.getWorkloads()['origin\u0000ollama\u0000\u0000upsert']).toBeUndefined() }) it('updates bridgeEpoch and invokes broker:resync plus each snapshot getter once when the epoch changes', async () => { diff --git a/desktop/tests/modular/inference-metrics.test.ts b/desktop/tests/modular/inference-metrics.test.ts index bf84ae730..b43981624 100644 --- a/desktop/tests/modular/inference-metrics.test.ts +++ b/desktop/tests/modular/inference-metrics.test.ts @@ -24,7 +24,7 @@ import { } from '@/ui/utils/inference-metrics' import type { InferenceMetricsTotals, NodeInferenceMetricsSnapshot } from '@/shared/types/metrics' import type { Workload, WorkloadLiveProgress } from '@/shared/types/workloads' -import { liveProgressKey, mergeWorkloadLiveProgress, workloadKey } from '@/shared/utils/workloads' +import { liveProgressKey, mergeWorkloadLiveProgress } from '@/shared/utils/workloads' vi.mock('electron', () => ({ BrowserWindow: { getAllWindows: () => [] } })) vi.mock('@/electron/window', () => ({ createOverviewWindow: vi.fn() })) @@ -1167,7 +1167,7 @@ describe('the workloads store keeps the live view of a running job', () => { ) // A live frame is a snapshot of a running request, not a catalog // event: it never enters the catalog. - expect(state.workloads.has(workloadKey('node-a', 'w1'))).toBe(false) + expect(state.workloads.size).toBe(0) }) it('keys the live view by the origin and the job id', async () => { diff --git a/desktop/tests/modular/uuid-node-keying.test.ts b/desktop/tests/modular/uuid-node-keying.test.ts index 5632b53f0..6b68c994e 100644 --- a/desktop/tests/modular/uuid-node-keying.test.ts +++ b/desktop/tests/modular/uuid-node-keying.test.ts @@ -206,8 +206,8 @@ describe('UUID node keying', () => { }) ) - const liveKey = 'uuid-wl-seed\u0000job-live' - const newKey = 'uuid-wl-seed\u0000job-new' + const liveKey = 'uuid-wl-seed\u0000ollama\u0000\u0000job-live' + const newKey = 'uuid-wl-seed\u0000ollama\u0000\u0000job-new' // The live entry is preserved, not clobbered by the older baseline row. expect(seeded[liveKey]).toMatchObject({ state: 'running', model: 'live-model' }) // A baseline job the stream had not delivered is filled in. diff --git a/desktop/tests/modular/workload-recovery.test.ts b/desktop/tests/modular/workload-recovery.test.ts index 9ae949b16..29137775b 100644 --- a/desktop/tests/modular/workload-recovery.test.ts +++ b/desktop/tests/modular/workload-recovery.test.ts @@ -7,11 +7,30 @@ import { parseWorkloadsInitial } from '@/electron/service-bridge/modular-state' import { emitBridgePush } from '@/electron/service-bridge/broadcaster' +import { workloadMatchesRemoval } from '@/shared/utils/workloads' vi.mock('@/electron/service-bridge/broadcaster', () => ({ emitBridgePush: vi.fn() })) const state = getModularBridgeState() +function key(originatedFrom: string, engine: string, id: string, runId = ''): string { + return [originatedFrom, engine, runId, id].join('\u0000') +} + +function upsert(id: string, engine: string, runId: string | null = null): void { + state.upsertWorkloadFromInfo({ + workloadInfo: { + id, + engine, + runId, + state: 'running', + model: engine, + originatedFrom: 'origin', + createdAt: 1 + } + }) +} + afterEach(() => { state.clearWorkloads() vi.clearAllMocks() @@ -21,7 +40,6 @@ describe('epoch workload snapshots', () => { it('releases modification tracking on cancellation and never lets an old token end a new fetch', () => { const old = state.beginWorkloadRecovery() state.removeWorkloadFromParams({ workloadId: 'removed', originatedFrom: 'origin' }) - expect(old.size).toBe(1) state.endWorkloadRecovery(old) expect(old.size).toBe(0) @@ -38,7 +56,9 @@ describe('epoch workload snapshots', () => { }) state.endWorkloadRecovery(old, []) state.endWorkloadRecovery(fresh, []) - expect(state.getWorkloads()['origin\u0000live']).toMatchObject({ model: 'live' }) + expect(state.getWorkloads()[key('origin', 'ollama', 'live')]).toMatchObject({ + model: 'live' + }) expect(fresh.size).toBe(0) const next = state.beginWorkloadRecovery() @@ -57,10 +77,12 @@ describe('epoch workload snapshots', () => { ] }) ) - expect(Object.keys(state.getWorkloads())).toEqual(['origin\u0000removed']) + expect(Object.keys(state.getWorkloads())).toEqual([key('origin', 'ollama', 'removed')]) expect(emitBridgePush).toHaveBeenCalledWith('workloads:remove', { workloadId: 'live', - originatedFrom: 'origin' + originatedFrom: 'origin', + engine: 'ollama', + runId: null }) expect(emitBridgePush).toHaveBeenCalledWith( 'workloads:upsert', @@ -68,3 +90,106 @@ describe('epoch workload snapshots', () => { ) }) }) + +describe('job identity', () => { + // Every proxy numbers its jobs from 1, so an Ollama job "3" and an oMLX + // job "3" on one node are two jobs, and so are a job and the one a + // restarted proxy gave the same number. + it('keeps jobs that share a number on different engines and runs apart', () => { + upsert('3', 'ollama') + upsert('3', 'omlx') + upsert('3', 'omlx', 'run-2') + expect(Object.keys(state.getWorkloads()).sort()).toEqual( + [ + key('origin', 'ollama', '3'), + key('origin', 'omlx', '3'), + key('origin', 'omlx', '3', 'run-2') + ].sort() + ) + }) + + // The broker's removal names only (originatedFrom, id) and drops every + // matching job, as the broker's own store does. + it('drops every job a broker removal names', () => { + upsert('3', 'ollama') + upsert('3', 'omlx') + upsert('4', 'omlx') + state.removeWorkloadFromParams({ workloadId: '3', originatedFrom: 'origin' }) + expect(Object.keys(state.getWorkloads())).toEqual([key('origin', 'omlx', '4')]) + }) + + // A removal seen while a snapshot is fetched also retires a matching job + // that only the (older) snapshot holds. + it('does not let a recovery snapshot bring back a job removed during the fetch', () => { + const snapshot = state.beginWorkloadRecovery() + state.removeWorkloadFromParams({ workloadId: '3', originatedFrom: 'origin' }) + state.endWorkloadRecovery( + snapshot, + parseWorkloadsInitial({ + workloads: [ + { + id: '3', + engine: 'omlx', + state: 'running', + model: 'm', + originatedFrom: 'origin', + createdAt: 1 + }, + { + id: '4', + engine: 'omlx', + state: 'running', + model: 'm', + originatedFrom: 'origin', + createdAt: 1 + } + ] + }) + ) + expect(Object.keys(state.getWorkloads())).toEqual([key('origin', 'omlx', '4')]) + }) +}) + +describe('exact removals', () => { + // A removal the desktop raises names the job's run, and "no run" is a run + // of its own: it must not take a same-numbered job from another run. + it('treats runId null as the job with no run, not any run', () => { + upsert('3', 'omlx') + upsert('3', 'omlx', 'run-2') + const snapshot = state.beginWorkloadRecovery() + state.endWorkloadRecovery( + snapshot, + parseWorkloadsInitial({ + workloads: [ + { + id: '3', + engine: 'omlx', + runId: 'run-2', + state: 'running', + model: 'm', + originatedFrom: 'origin', + createdAt: 1 + } + ] + }) + ) + expect(emitBridgePush).toHaveBeenCalledWith('workloads:remove', { + workloadId: '3', + originatedFrom: 'origin', + engine: 'omlx', + runId: null + }) + expect( + workloadMatchesRemoval( + { id: '3', originatedFrom: 'origin', engine: 'omlx', runId: 'run-2' }, + { + workloadId: '3', + originatedFrom: 'origin', + engine: 'omlx', + runId: null + } + ) + ).toBe(false) + expect(Object.keys(state.getWorkloads())).toEqual([key('origin', 'omlx', '3', 'run-2')]) + }) +}) diff --git a/services/nvpair-tui/ui/model.go b/services/nvpair-tui/ui/model.go index 55106bff0..2e2109b1b 100644 --- a/services/nvpair-tui/ui/model.go +++ b/services/nvpair-tui/ui/model.go @@ -54,6 +54,12 @@ type Model struct { recoveryGeneration uint64 recoveryActive bool recoveryPushes map[string]NotificationMsg + // recoveryWorkloads holds workload upserts and removals seen during + // recovery, in arrival order. They cannot share recoveryPushes: a removal + // names only (origin, id) and drops every matching job, so the order of a + // removal and a later upsert decides the outcome, and a map replays in no + // particular order. + recoveryWorkloads []NotificationMsg } // New builds the root model over a connected broker client, the broker's @@ -157,7 +163,9 @@ func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.brokerVersion = readyVersion(msg.Msg) } if m.recoveryActive { - if key := recoveryPushKey(msg); key != "" { + if isWorkloadPush(msg.Msg.Method) { + m.recoveryWorkloads = append(m.recoveryWorkloads, msg) + } else if key := recoveryPushKey(msg); key != "" { m.recoveryPushes[key] = msg } } @@ -201,7 +209,11 @@ func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { } cmds = append(cmds, m.broadcast(push)...) } + for _, push := range m.recoveryWorkloads { + cmds = append(cmds, m.broadcast(push)...) + } m.recoveryPushes = nil + m.recoveryWorkloads = nil // Hydrated even if optional getters errored: their views received the // error result, and healthy same-epoch traffic must not re-burst (§4). m.hydratedEpoch = msg.epoch diff --git a/services/nvpair-tui/ui/model_test.go b/services/nvpair-tui/ui/model_test.go index 07e1f1a29..bb82b0f7f 100644 --- a/services/nvpair-tui/ui/model_test.go +++ b/services/nvpair-tui/ui/model_test.go @@ -381,16 +381,55 @@ func TestRecoveryPreservesConcurrentSnapshotPushesAndWorkloadTombstones(t *testi b.push(t, "A", 3, "workloads:upsert", map[string]any{"workloadInfo": workload{ID: "removed", OriginatedFrom: "x"}}) b.push(t, "A", 4, "workloads:remove", map[string]string{"workloadId": "removed", "originatedFrom": "x"}) b.push(t, "A", 5, "workloads:upsert", map[string]any{"workloadInfo": workload{ID: "live", OriginatedFrom: "x"}}) - l.until(t, func() bool { return len(l.m.recoveryPushes) == 3 }) + l.until(t, func() bool { return len(l.m.recoveryPushes) == 1 && len(l.m.recoveryWorkloads) == 3 }) b.unhold() l.until(t, func() bool { return l.m.hydratedEpoch == "A" }) n := l.m.views[2].(*nodesView) w := l.m.views[4].(*workloadsView) - if len(n.nodes) != 1 || n.nodes[0].HostUUID != "live" || len(w.byKey) != 1 || w.byKey[workloadKey("x", "live")].ID != "live" { + if len(n.nodes) != 1 || n.nodes[0].HostUUID != "live" || len(w.byKey) != 1 || w.byKey[workloadKey(workload{ID: "live", OriginatedFrom: "x"})].ID != "live" { t.Fatalf("newer pushes overwritten: %v %v", n.nodes, w.byKey) } } +// Jobs that share a number on different engines are separate jobs, and a +// removal names only (origin, id). Replayed in order, the journal keeps both +// engines' upserts, and an upsert that follows a removal survives it while the +// other engine's job stays removed. +func TestRecoveryKeepsSameNumberedJobsAndReplaysInOrder(t *testing.T) { + b, l := newRecoverySession(t) + b.mu.Lock() + b.workloads = []workload{{ID: "3", Engine: "ollama", OriginatedFrom: "x"}, {ID: "3", Engine: "omlx", OriginatedFrom: "x"}} + b.mu.Unlock() + entry := b.hold("errors:get-initial") + b.push(t, "A", 1, "heartbeat", nil) + l.until(t, func() bool { return l.m.recoveryActive }) + awaitEntry(t, entry) + b.push(t, "A", 2, "workloads:upsert", map[string]any{"workloadInfo": workload{ID: "7", Engine: "ollama", OriginatedFrom: "x"}}) + b.push(t, "A", 3, "workloads:upsert", map[string]any{"workloadInfo": workload{ID: "7", Engine: "omlx", OriginatedFrom: "x"}}) + b.push(t, "A", 4, "workloads:remove", map[string]string{"workloadId": "3", "originatedFrom": "x"}) + b.push(t, "A", 5, "workloads:upsert", map[string]any{"workloadInfo": workload{ID: "3", Engine: "omlx", OriginatedFrom: "x", State: "running"}}) + l.until(t, func() bool { return len(l.m.recoveryWorkloads) == 4 }) + b.unhold() + l.until(t, func() bool { return l.m.hydratedEpoch == "A" }) + w := l.m.views[4].(*workloadsView) + want := []workload{ + {ID: "7", Engine: "ollama", OriginatedFrom: "x"}, + {ID: "7", Engine: "omlx", OriginatedFrom: "x"}, + {ID: "3", Engine: "omlx", OriginatedFrom: "x"}, + } + if len(w.byKey) != len(want) { + t.Fatalf("rows = %v, want %d", w.byKey, len(want)) + } + for _, j := range want { + if _, ok := w.byKey[workloadKey(j)]; !ok { + t.Fatalf("missing %s/%s in %v", j.Engine, j.ID, w.byKey) + } + } + if w.byKey[workloadKey(workload{ID: "3", Engine: "omlx", OriginatedFrom: "x"})].State != "running" { + t.Fatal("the upsert after the removal did not win") + } +} + // An unavailable optional worker is a steady state, not a reason to re-burst: // every domain whose getter succeeded gets its baseline, the epoch counts as // hydrated, and healthy same-epoch traffic issues no further RPCs at all. diff --git a/services/nvpair-tui/ui/recovery.go b/services/nvpair-tui/ui/recovery.go index 3e890bec2..5f0e3086a 100644 --- a/services/nvpair-tui/ui/recovery.go +++ b/services/nvpair-tui/ui/recovery.go @@ -75,6 +75,7 @@ func (m *Model) beginRecovery() tea.Cmd { m.recoveryActive = true m.recoveryGeneration++ m.recoveryPushes = map[string]NotificationMsg{} + m.recoveryWorkloads = nil generation, epoch := m.recoveryGeneration, m.epoch client := m.client reads := []recoveryRead{ @@ -187,22 +188,15 @@ func recoveryPushKey(msg NotificationMsg) string { var e engineStatus _ = decodeParams(msg.Msg.Params, &e) return method + ":" + e.Engine - case "workloads:upsert": - var p struct { - WorkloadInfo workload `json:"workloadInfo"` - } - _ = decodeParams(msg.Msg.Params, &p) - return "workload:" + workloadKey(p.WorkloadInfo.OriginatedFrom, p.WorkloadInfo.ID) - case "workloads:remove": - var p struct { - WorkloadID string `json:"workloadId"` - OriginatedFrom string `json:"originatedFrom"` - } - _ = decodeParams(msg.Msg.Params, &p) - return "workload:" + workloadKey(p.OriginatedFrom, p.WorkloadID) } if proxyPush(method) { return method } return "" } + +// isWorkloadPush reports whether a notification changes the workload list. +// Those are journaled in order during recovery rather than by identity. +func isWorkloadPush(method string) bool { + return method == "workloads:upsert" || method == "workloads:remove" +} diff --git a/services/nvpair-tui/ui/workloads.go b/services/nvpair-tui/ui/workloads.go index 5b19ce9e5..76150d58c 100644 --- a/services/nvpair-tui/ui/workloads.go +++ b/services/nvpair-tui/ui/workloads.go @@ -18,6 +18,7 @@ type workload struct { Engine string `json:"engine"` State string `json:"state"` OriginatedFrom string `json:"originatedFrom"` + RunID string `json:"runId"` ScheduledOn string `json:"scheduledOn"` CreatedAt int64 `json:"createdAt"` // Unix millis } @@ -140,7 +141,7 @@ func (v *workloadsView) Update(msg tea.Msg) tea.Cmd { OriginatedFrom string `json:"originatedFrom"` } _ = decodeParams(msg.Msg.Params, &p) - v.remove(workloadKey(p.OriginatedFrom, p.WorkloadID)) + v.remove(p.OriginatedFrom, p.WorkloadID) } return nil @@ -153,7 +154,7 @@ func (v *workloadsView) Update(msg tea.Msg) tea.Cmd { } func (v *workloadsView) upsert(w workload) { - key := workloadKey(w.OriginatedFrom, w.ID) + key := workloadKey(w) if _, ok := v.byKey[key]; !ok { v.order = append(v.order, key) } @@ -161,18 +162,24 @@ func (v *workloadsView) upsert(w workload) { v.refreshRows() } -func (v *workloadsView) remove(key string) { - if _, ok := v.byKey[key]; !ok { - return - } - delete(v.byKey, key) - for i, k := range v.order { - if k == key { - v.order = append(v.order[:i], v.order[i+1:]...) - break +// remove drops every row a workloads:remove names. The removal carries only +// (originatedFrom, id), so, like the broker's own store, it drops the job on +// every engine and run of that origin. +func (v *workloadsView) remove(origin, id string) { + kept := v.order[:0] + removed := false + for _, k := range v.order { + if w := v.byKey[k]; w.OriginatedFrom == origin && w.ID == id { + delete(v.byKey, k) + removed = true + continue } + kept = append(kept, k) + } + v.order = kept + if removed { + v.refreshRows() } - v.refreshRows() } func (v *workloadsView) refreshRows() { @@ -210,4 +217,9 @@ func (v *workloadsView) View() string { func (v *workloadsView) Help() []key.Binding { return nil } -func workloadKey(origin, id string) string { return origin + "/" + id } +// workloadKey is the backend's own job identity. Every proxy numbers its jobs +// from 1, and again from 1 when it restarts, so the origin, engine and run are +// all needed to keep jobs that share a number apart. +func workloadKey(w workload) string { + return w.OriginatedFrom + "/" + w.Engine + "/" + w.RunID + "/" + w.ID +} diff --git a/services/nvpair-tui/ui/workloads_identity_test.go b/services/nvpair-tui/ui/workloads_identity_test.go new file mode 100644 index 000000000..8b27c34bb --- /dev/null +++ b/services/nvpair-tui/ui/workloads_identity_test.go @@ -0,0 +1,24 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package ui + +import "testing" + +// Every proxy numbers its jobs from 1, so jobs that share a number on different +// engines or runs are different jobs, and a removal (which names only the +// origin and id) drops all of them. +func TestWorkloadsKeepJobsThatShareANumberApart(t *testing.T) { + v := &workloadsView{byKey: map[string]workload{}} + v.upsert(workload{ID: "3", Engine: "ollama", OriginatedFrom: "x"}) + v.upsert(workload{ID: "3", Engine: "omlx", OriginatedFrom: "x"}) + v.upsert(workload{ID: "3", Engine: "omlx", RunID: "run-2", OriginatedFrom: "x"}) + v.upsert(workload{ID: "4", Engine: "omlx", OriginatedFrom: "x"}) + if len(v.byKey) != 4 || len(v.order) != 4 { + t.Fatalf("rows = %d (order %d), want 4", len(v.byKey), len(v.order)) + } + v.remove("x", "3") + if len(v.byKey) != 1 || len(v.order) != 1 || v.byKey[v.order[0]].ID != "4" { + t.Fatalf("after removing job 3: %v", v.byKey) + } +} diff --git a/services/nvpair-workload-manager/cluster_mtls_test.go b/services/nvpair-workload-manager/cluster_mtls_test.go index cd2256687..b938bfbb0 100644 --- a/services/nvpair-workload-manager/cluster_mtls_test.go +++ b/services/nvpair-workload-manager/cluster_mtls_test.go @@ -81,6 +81,17 @@ func setupNode(t *testing.T, certPEM, keyPEM []byte, pins map[string][]byte) str return dir } +// newPinnedPeerDirs builds two mutually pinned nodes. Tests that construct a +// Manager need selfDir because NewManager opens its own mesh from that path. +func newPinnedPeerDirs(t *testing.T) (selfDir, peerDir string) { + t.Helper() + selfCert, selfKey := genLeaf(t, "uuid-self") + peerCert, peerKey := genLeaf(t, "uuid-peer") + selfDir = setupNode(t, selfCert, selfKey, map[string][]byte{"uuid-peer": peerCert}) + peerDir = setupNode(t, peerCert, peerKey, map[string][]byte{"uuid-self": selfCert}) + return selfDir, peerDir +} + // newPinnedPeerMeshes builds the two sides of a minimal two-node cluster: self // pins peer and peer pins self, so a test can drive the real inter-node path // (cluster mTLS from a pinned caller) rather than a plaintext shortcut. There is @@ -88,10 +99,7 @@ func setupNode(t *testing.T, certPEM, keyPEM []byte, pins map[string][]byte) str // receiver test goes through here. func newPinnedPeerMeshes(t *testing.T) (self, peer *clustertrust.Mesh) { t.Helper() - selfCert, selfKey := genLeaf(t, "uuid-self") - peerCert, peerKey := genLeaf(t, "uuid-peer") - selfDir := setupNode(t, selfCert, selfKey, map[string][]byte{"uuid-peer": peerCert}) - peerDir := setupNode(t, peerCert, peerKey, map[string][]byte{"uuid-self": selfCert}) + selfDir, peerDir := newPinnedPeerDirs(t) self, peer = clustertrust.Open(selfDir), clustertrust.Open(peerDir) if !self.Clustered() || !peer.Clustered() { t.Fatal("a dir holding a keypair and a pin must read as clustered") diff --git a/services/nvpair-workload-manager/dedup.go b/services/nvpair-workload-manager/dedup.go index c265a6221..0012bbb97 100644 --- a/services/nvpair-workload-manager/dedup.go +++ b/services/nvpair-workload-manager/dedup.go @@ -12,10 +12,10 @@ import ( // session-scoped volume at ~dozen-node scale (spec §5). const defaultDedupCapacity = 10000 -// dedupIndex is a bounded LRU set of keys. It answers a single question: -// "have I seen this key before?" and records it if not, evicting the -// least-recently-seen key once capacity is exceeded. It is safe for -// concurrent use — the inter-node HTTP handler runs one goroutine per +// dedupIndex is a bounded LRU set of successfully emitted keys. It records a +// key only after the broker emit succeeds, and serializes concurrent emits for +// the same key. Completed keys are evicted least-recently-used first. It is +// safe for concurrent use — the inter-node HTTP handler runs one goroutine per // request. // // Keys are opaque strings built by the caller: lifecycle events key on @@ -30,6 +30,7 @@ type dedupIndex struct { capacity int ll *list.List // front = most recently seen items map[string]*list.Element // key -> element in ll + inFlight map[string]chan struct{} // one broker emit at a time per key } func newDedupIndex(capacity int) *dedupIndex { @@ -40,21 +41,54 @@ func newDedupIndex(capacity int) *dedupIndex { capacity: capacity, ll: list.New(), items: make(map[string]*list.Element, capacity), + inFlight: make(map[string]chan struct{}), } } -// seenOrAdd returns true if the key was already present (a duplicate). On a -// first sighting it records the key and returns false. Either way the key is -// promoted to most-recently-seen. -func (d *dedupIndex) seenOrAdd(key string) bool { - d.mu.Lock() - defer d.mu.Unlock() +// emitOnce serializes the check, broker emit, and record for one key. Other +// keys can emit concurrently. If emit fails, the key stays absent and a waiting +// request can retry it. A completed key is reported as a duplicate. +func (d *dedupIndex) emitOnce(key string, emit func() error) (bool, error) { + for { + d.mu.Lock() + if el, ok := d.items[key]; ok { + d.ll.MoveToFront(el) + d.mu.Unlock() + return true, nil + } + if done, ok := d.inFlight[key]; ok { + d.mu.Unlock() + <-done + continue + } + done := make(chan struct{}) + d.inFlight[key] = done + d.mu.Unlock() - if el, ok := d.items[key]; ok { - d.ll.MoveToFront(el) - return true + // Release waiters even if an emitter panics and net/http recovers the + // request. A panicking emit has not completed successfully. + err := func() (err error) { + completed := false + defer func() { + d.mu.Lock() + if completed && err == nil { + d.addLocked(key) + } + delete(d.inFlight, key) + close(done) + d.mu.Unlock() + }() + err = emit() + completed = true + return err + }() + return false, err } +} +// addLocked records a new key and evicts the least-recently-seen key past +// capacity. The caller holds d.mu and has checked that key is absent. +func (d *dedupIndex) addLocked(key string) { el := d.ll.PushFront(key) d.items[key] = el if d.ll.Len() > d.capacity { @@ -64,7 +98,6 @@ func (d *dedupIndex) seenOrAdd(key string) bool { delete(d.items, oldest.Value.(string)) } } - return false } // keyLifecycle builds the dedup key for a lifecycle event. Workload.id is only diff --git a/services/nvpair-workload-manager/dedup_test.go b/services/nvpair-workload-manager/dedup_test.go index 557f54d70..39abdc05b 100644 --- a/services/nvpair-workload-manager/dedup_test.go +++ b/services/nvpair-workload-manager/dedup_test.go @@ -3,18 +3,181 @@ package main -import "testing" +import ( + "errors" + "sync/atomic" + "testing" + "time" +) -func TestDedupSeenOrAdd(t *testing.T) { +// emitDedupKey exercises the production dedup path with a successful broker +// emit, then reports whether that key was already recorded. +func emitDedupKey(t *testing.T, d *dedupIndex, key string) bool { + t.Helper() + duplicate, err := d.emitOnce(key, func() error { return nil }) + if err != nil { + t.Fatalf("emit dedup key: %v", err) + } + return duplicate +} + +// TestDedupIndex_EmitOnceSerializesOnlyMatchingKeys proves that an in-flight +// emit blocks only requests with the same key. The test controls this order: +// +// 1. The "first" goroutine reserves its key, enters its emit callback, and +// signals firstEntered. It then waits on releaseFirst, keeping that key +// in flight. +// 2. While "first" is still waiting, the "second" goroutine emits a +// different key. The test requires secondDone before closing releaseFirst. +// 3. The test closes releaseFirst and waits for the original emit to finish. +// +// A single lock held across all emit callbacks would make step 2 time out. +func TestDedupIndex_EmitOnceSerializesOnlyMatchingKeys(t *testing.T) { + d := newDedupIndex(10) + firstEntered := make(chan struct{}) + releaseFirst := make(chan struct{}) + // Release the blocked goroutine even if an assertion fails early. + defer func() { + select { + case <-releaseFirst: + default: + close(releaseFirst) + } + }() + firstDone := make(chan error, 1) + go func() { + _, err := d.emitOnce("first", func() error { + close(firstEntered) + <-releaseFirst + return nil + }) + firstDone <- err + }() + select { + case <-firstEntered: + case <-time.After(2 * time.Second): + t.Fatal("first key did not begin emitting") + } + + type result struct { + duplicate bool + err error + } + secondDone := make(chan result, 1) + go func() { + duplicate, err := d.emitOnce("second", func() error { return nil }) + secondDone <- result{duplicate: duplicate, err: err} + }() + select { + case got := <-secondDone: + if got.err != nil || got.duplicate { + t.Fatalf("different key result = (%t, %v), want (false, nil)", got.duplicate, got.err) + } + case <-time.After(2 * time.Second): + t.Fatal("different key blocked behind the first emit") + } + close(releaseFirst) + select { + case err := <-firstDone: + if err != nil { + t.Fatalf("first key emit: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("first key did not finish") + } +} + +// TestDedupIndex_WaiterRetriesAfterFailedEmit proves that a failed emit does +// not permanently reserve or record its key. Unlike the different-key test +// above, both goroutines use "same": +// +// 1. The first goroutine reserves "same", enters its emit callback, signals +// firstEntered, and waits on releaseFirst. +// 2. The second goroutine requests "same". It must wait while the first +// callback is in flight; returning or emitting before releaseFirst closes +// fails the test. +// 3. The test closes releaseFirst. The first callback returns an error, so +// "same" remains unrecorded. The waiting goroutine then emits it once and +// succeeds. +// +// The final retry count distinguishes a real retry from a duplicate response. +func TestDedupIndex_WaiterRetriesAfterFailedEmit(t *testing.T) { + d := newDedupIndex(10) + firstEntered := make(chan struct{}) + releaseFirst := make(chan struct{}) + defer func() { + select { + case <-releaseFirst: + default: + close(releaseFirst) + } + }() + wantErr := errors.New("broker unavailable") + firstDone := make(chan error, 1) + go func() { + _, err := d.emitOnce("same", func() error { + close(firstEntered) + <-releaseFirst + return wantErr + }) + firstDone <- err + }() + select { + case <-firstEntered: + case <-time.After(2 * time.Second): + t.Fatal("first request did not begin emitting") + } + + var retries atomic.Int32 + secondDone := make(chan error, 1) + go func() { + duplicate, err := d.emitOnce("same", func() error { + retries.Add(1) + return nil + }) + if duplicate { + secondDone <- errors.New("failed key was treated as a duplicate") + return + } + secondDone <- err + }() + select { + case err := <-secondDone: + t.Fatalf("waiting request completed before first emit: %v", err) + case <-time.After(50 * time.Millisecond): + } + close(releaseFirst) + select { + case err := <-firstDone: + if !errors.Is(err, wantErr) { + t.Fatalf("first emit error = %v, want %v", err, wantErr) + } + case <-time.After(2 * time.Second): + t.Fatal("first request did not finish") + } + select { + case err := <-secondDone: + if err != nil { + t.Fatalf("waiting retry: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("waiting retry did not finish") + } + if got := retries.Load(); got != 1 { + t.Fatalf("retry emits = %d, want 1", got) + } +} + +func TestDedupEmitOnce(t *testing.T) { d := newDedupIndex(8) - if d.seenOrAdd("a") { + if emitDedupKey(t, d, "a") { t.Fatal("first sighting of a should be new") } - if !d.seenOrAdd("a") { + if !emitDedupKey(t, d, "a") { t.Fatal("second sighting of a should be a duplicate") } - if d.seenOrAdd("b") { + if emitDedupKey(t, d, "b") { t.Fatal("first sighting of b should be new") } } @@ -22,11 +185,11 @@ func TestDedupSeenOrAdd(t *testing.T) { func TestDedupEviction(t *testing.T) { d := newDedupIndex(2) - d.seenOrAdd("a") // {a} - d.seenOrAdd("b") // {a,b} - d.seenOrAdd("c") // evicts a -> {b,c} + emitDedupKey(t, d, "a") // {a} + emitDedupKey(t, d, "b") // {a,b} + emitDedupKey(t, d, "c") // evicts a -> {b,c} - if d.seenOrAdd("a") { + if emitDedupKey(t, d, "a") { t.Fatal("a should have been evicted and read as new") } } @@ -37,17 +200,17 @@ func TestDedupKeysDistinguishStateAndKind(t *testing.T) { w := &Workload{ID: "wl-1", OriginatedFrom: "node-A", State: StateQueued} wRunning := &Workload{ID: "wl-1", OriginatedFrom: "node-A", State: StateRunning} - if d.seenOrAdd(keyLifecycle(w)) { + if emitDedupKey(t, d, keyLifecycle(w)) { t.Fatal("queued should be new") } - if d.seenOrAdd(keyLifecycle(wRunning)) { + if emitDedupKey(t, d, keyLifecycle(wRunning)) { t.Fatal("running for same id is a different key, should be new") } - if !d.seenOrAdd(keyLifecycle(w)) { + if !emitDedupKey(t, d, keyLifecycle(w)) { t.Fatal("repeat queued should dedup") } // A removal keyed on the same id must not collide with a lifecycle key. - if d.seenOrAdd(keyRemove("node-A", "wl-1")) { + if emitDedupKey(t, d, keyRemove("node-A", "wl-1")) { t.Fatal("removal of wl-1 must not collide with lifecycle keys") } } @@ -62,13 +225,13 @@ func TestDedupDistinguishesNodes(t *testing.T) { nodeA := &Workload{ID: "wl-1", OriginatedFrom: "node-A", State: StateQueued} nodeB := &Workload{ID: "wl-1", OriginatedFrom: "node-B", State: StateQueued} - if d.seenOrAdd(keyLifecycle(nodeA)) { + if emitDedupKey(t, d, keyLifecycle(nodeA)) { t.Fatal("node-A wl-1 should be new") } - if d.seenOrAdd(keyLifecycle(nodeB)) { + if emitDedupKey(t, d, keyLifecycle(nodeB)) { t.Fatal("node-B wl-1 has the same id but a different node, must not dedup against node-A") } - if !d.seenOrAdd(keyLifecycle(nodeA)) { + if !emitDedupKey(t, d, keyLifecycle(nodeA)) { t.Fatal("repeat of node-A wl-1 should dedup") } } @@ -82,16 +245,16 @@ func TestDedupDistinguishesEngineAndRun(t *testing.T) { lmstudio := &Workload{ID: "1", OriginatedFrom: "host", Engine: "lmstudio", RunID: "r2", State: StateRunning} restarted := &Workload{ID: "1", OriginatedFrom: "host", Engine: "ollama", RunID: "r3", State: StateRunning} - if d.seenOrAdd(keyLifecycle(ollama)) { + if emitDedupKey(t, d, keyLifecycle(ollama)) { t.Fatal("ollama host/1 should be new") } - if d.seenOrAdd(keyLifecycle(lmstudio)) { + if emitDedupKey(t, d, keyLifecycle(lmstudio)) { t.Fatal("lmstudio host/1 shares the id but a different engine; must not dedup") } - if d.seenOrAdd(keyLifecycle(restarted)) { + if emitDedupKey(t, d, keyLifecycle(restarted)) { t.Fatal("a reused id from a new run must not dedup against the old run") } - if !d.seenOrAdd(keyLifecycle(ollama)) { + if !emitDedupKey(t, d, keyLifecycle(ollama)) { t.Fatal("repeat of ollama host/1 should dedup") } } @@ -102,13 +265,13 @@ func TestDedupDistinguishesEngineAndRun(t *testing.T) { func TestDedupRemovalDistinguishesNodes(t *testing.T) { d := newDedupIndex(8) - if d.seenOrAdd(keyRemove("node-A", "wl-1")) { + if emitDedupKey(t, d, keyRemove("node-A", "wl-1")) { t.Fatal("removal of node-A wl-1 should be new") } - if d.seenOrAdd(keyRemove("node-B", "wl-1")) { + if emitDedupKey(t, d, keyRemove("node-B", "wl-1")) { t.Fatal("removal of node-B wl-1 must not dedup against node-A") } - if !d.seenOrAdd(keyRemove("node-A", "wl-1")) { + if !emitDedupKey(t, d, keyRemove("node-A", "wl-1")) { t.Fatal("repeat removal of node-A wl-1 should dedup") } } diff --git a/services/nvpair-workload-manager/manager.go b/services/nvpair-workload-manager/manager.go index 25fe126f9..fe066ad1e 100644 --- a/services/nvpair-workload-manager/manager.go +++ b/services/nvpair-workload-manager/manager.go @@ -27,10 +27,11 @@ import ( const discoveryInterval = 5 * time.Second // resyncInterval is how often each node re-asserts its own active + recently- -// terminal workloads to peers (the anti-entropy heartbeat), so a dropped -// delivery or a peer's wrong node-loss guess reconciles within a couple of -// intervals. terminalRetention keeps a finished workload in the re-sync set for -// two intervals — so it is re-asserted ~twice before ageing out. +// terminal workloads to peers (the anti-entropy heartbeat). A later +// re-assertion can repair a missed lifecycle event or a peer's wrong node-loss +// guess if it is queued and delivered. terminalRetention keeps a finished +// workload in the re-sync set for two intervals, allowing roughly two +// re-assertions before it ages out. Removals are not kept in that set. // // resyncInterval is load-bearing OUTSIDE this binary. nvpair-ui-broker's // staleness sweep (workloadOriginSilenceTimeout in its broker.go) retires a @@ -59,9 +60,9 @@ type workloadKey struct { // workloadEvent is a stored lifecycle notification (method + params) for a // workloadKey. A terminal event is retained until expiresAt so the heartbeat -// re-asserts it a couple of times (covering a dropped delivery / a peer's wrong -// node-loss guess); an active event has a zero expiresAt and is retained until -// it terminates or is removed. +// can re-assert it a couple of times (offering another chance after a missed +// delivery or a peer's wrong node-loss guess); an active event has a zero +// expiresAt and is retained until it terminates or is removed. type workloadEvent struct { method string params json.RawMessage @@ -90,11 +91,23 @@ type Manager struct { // activeLocal is this node's re-sync set: the latest event per local-origin // (origin,id) — active workloads plus recently-terminal ones (retained // until they expire). It backfills a newly-discovered peer (pushActiveSnapshot) - // and is re-asserted on the heartbeat (resyncLoop) so peers reconcile to this - // node's authoritative state. Guarded by activeMu. + // and is re-asserted on the heartbeat (resyncLoop) to help peers reconcile + // to this node's authoritative state. Guarded by activeMu. activeMu sync.Mutex activeLocal map[workloadKey]workloadEvent + // broadcastMu keeps snapshot enqueueing ordered with local lifecycle and + // removal updates. Always acquire it before activeMu. + broadcastMu sync.Mutex + + // broadcastCh serializes outbound inter-node frames in the order the + // read loop produced them. broadcastFrame only enqueues (never blocks on + // network I/O), and a single worker drains the queue in order — so a + // queued remove cannot overtake the lifecycle upsert it follows. Without + // this, each frame fanned out in its own goroutine and a late upsert + // could resurrect a workload on peers that had already removed it. + broadcastCh chan []byte + ctx context.Context cancel context.CancelFunc } @@ -127,6 +140,7 @@ func NewManager(codec *Codec, port int, selfUUID, clusterDir string) *Manager { peerSource: relaySource, relaySource: relaySource, activeLocal: make(map[workloadKey]workloadEvent), + broadcastCh: make(chan []byte, broadcastQueueDepth), } m.server = NewServer(port, dedup, mesh, m.emitUpsert, m.emitRemove, m.emitProgress) return m @@ -160,6 +174,10 @@ func (m *Manager) Run(ctx context.Context) error { go m.discoveryLoop(ctx) go m.resyncLoop(ctx) + // The single ordered broadcast consumer: frames go out in the order the + // read loop produced them, so a queued remove cannot overtake the + // lifecycle event it follows. + go m.broadcastLoop(ctx) // Follow this node into and out of a cluster. Every gate already reads live // membership, so the watch exists to notice a change with no traffic flowing // and to re-assert our workloads immediately: peers that could not receive @@ -357,6 +375,8 @@ func (m *Manager) handleLocalLifecycle(msg *Message) { return } key := workloadKey{origin: wl.OriginatedFrom, engine: wl.Engine, runID: wl.RunID, id: wl.ID} + m.broadcastMu.Lock() + defer m.broadcastMu.Unlock() m.trackActive(key, msg.Method, msg.Params, wl.State) slog.Debug("broadcasting local lifecycle", "method", msg.Method, "id", wl.ID, "state", wl.State, "peers", m.peers.count()) m.broadcastFrame(msg.Method, msg.Params) @@ -390,24 +410,59 @@ func (m *Manager) handleLocalRemove(msg *Message) { // is preserved for peers' dedup. The removal wire carries only // (workloadId, originatedFrom) — no engine/runId — so drop every composite // key matching that pair. + m.broadcastMu.Lock() + defer m.broadcastMu.Unlock() m.untrackActive(nodeID, workloadID) slog.Debug("broadcasting local removal", "workloadId", workloadID, "node", nodeID, "peers", m.peers.count()) m.broadcastFrame(msg.Method, msg.Params) } -// broadcastFrame re-marshals a single notification and fans it out to every peer -// asynchronously, so a slow peer never blocks the read loop. Delivery is -// immediate and per-event (no batching/conflation): the origin's own view -// already updated synchronously in the broker, and peers must see each -// transition promptly and individually — a batching window would add latency and -// drop intermediate states, skewing each node's independent scheduling view. +// broadcastQueueDepth bounds how many outbound frames can wait for the +// ordered broadcast worker. This bounds queued memory without blocking the +// read loop: a full queue drops each new frame with a warning. Heartbeat and +// backfill can reassert tracked lifecycle state, but neither replays removals. +// If this is too small, we can make it bigger or configurable. +const broadcastQueueDepth = 1024 + +// broadcastFrame re-marshals a single notification and enqueues it for the +// ordered broadcast worker. Delivery is still asynchronous — a slow peer never +// blocks the read loop — but frames now go out in the order they were produced +// (no batching/conflation): the origin's own view already updated +// synchronously in the broker. Each queued frame is sent separately, but +// delivery remains best-effort and a full queue drops the new frame. func (m *Manager) broadcastFrame(method string, params json.RawMessage) { frame, err := json.Marshal(&Message{JSONRPC: "2.0", Method: method, Params: params}) if err != nil { slog.Error("failed to marshal broadcast frame", "method", method, "err", err) return } - go m.broadcaster.Broadcast(m.ctx, frame) + if m.broadcastCh == nil { + // Only reachable by hand-built Managers (tests); NewManager always + // installs the queue. + slog.Warn("broadcast queue not initialized, dropping frame", "method", method) + return + } + select { + case m.broadcastCh <- frame: + default: + slog.Warn("broadcast queue full, dropping frame", "method", method) + } +} + +// broadcastLoop is the single ordered consumer of broadcastCh. One worker (not +// a goroutine per frame) preserves queue order, so a queued remove cannot +// overtake the lifecycle event it follows. It does not guarantee delivery. +// Broadcast aborts in-flight attempts on ctx cancellation, so at shutdown the +// loop just exits; frames still queued are dropped with the process. +func (m *Manager) broadcastLoop(ctx context.Context) { + for { + select { + case <-ctx.Done(): + return + case frame := <-m.broadcastCh: + m.broadcaster.Broadcast(ctx, frame) + } + } } // trackActive records the latest event for a local-origin workload. A @@ -479,8 +534,9 @@ func (m *Manager) pushActiveSnapshot() { // resyncLoop is the anti-entropy heartbeat: every resyncInterval it re-asserts // this node's own active + recently-terminal workloads to all peers. Because the // origin is the single writer for its workloads, a peer that missed a delivery -// or made a wrong node-loss guess reconciles to the origin's authoritative state -// within a couple of intervals. +// or made a wrong node-loss guess may reconcile to the origin's authoritative +// state when later reassertions are queued and delivered. A full queue can drop +// those reassertions, and removals are not part of the re-sync set. func (m *Manager) resyncLoop(ctx context.Context) { ticker := time.NewTicker(resyncInterval) defer ticker.Stop() @@ -497,6 +553,8 @@ func (m *Manager) resyncLoop(ctx context.Context) { // broadcastSnapshot re-broadcasts the current re-sync set, one per-event frame // each. Shared by the discovery backfill and the heartbeat. func (m *Manager) broadcastSnapshot(reason string) { + m.broadcastMu.Lock() + defer m.broadcastMu.Unlock() snapshot := m.activeSnapshot() if len(snapshot) == 0 { return diff --git a/services/nvpair-workload-manager/manager_test.go b/services/nvpair-workload-manager/manager_test.go new file mode 100644 index 000000000..2198df5cc --- /dev/null +++ b/services/nvpair-workload-manager/manager_test.go @@ -0,0 +1,226 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "encoding/json" + "fmt" + "io" + "log/slog" + "net" + "net/http" + "net/http/httptest" + "strconv" + "sync" + "testing" + "time" + + "nvpair-shared/clustertrust" +) + +// codecNop is a throwaway io.ReadWriter for a Manager whose broker side is +// never exercised. +type codecNop struct{} + +func (codecNop) Read([]byte) (int, error) { return 0, io.EOF } +func (codecNop) Write(p []byte) (int, error) { return len(p), nil } + +type jSequence struct { + Seq int `json:"seq"` +} + +type jParams struct { + Params jSequence `json:"params"` +} + +func assertFrame(t *testing.T, body []byte, want int) { + t.Helper() + var frame jParams + if err := json.Unmarshal(body, &frame); err != nil { + t.Fatalf("decode frame %d: %v", want, err) + } + if frame.Params.Seq != want { + t.Fatalf("frame %d arrived with sequence %d", want, frame.Params.Seq) + } +} + +func newBroadcastManagerForPeer(t *testing.T, selfDir string, peer *httptest.Server) *Manager { + t.Helper() + host, portStr, err := net.SplitHostPort(peer.Listener.Addr().String()) + if err != nil { + t.Fatalf("split test peer address: %v", err) + } + port, err := strconv.Atoi(portStr) + if err != nil { + t.Fatalf("parse test peer port: %v", err) + } + + m := NewManager(NewCodec(codecNop{}), 0, "uuid-self", selfDir) + m.peers.Replace([]PeerNode{{ + ID: "peer-1", Addresses: []string{host}, Port: port, + TXT: []string{"cluster-uuid=uuid-peer"}, + }}) + return m +} + +// TestManager_BroadcastPreservesEnqueueOrder checks that frames queued in order +// reach a peer in the same order over the cluster-mTLS broadcast path. +func TestManager_BroadcastPreservesEnqueueOrder(t *testing.T) { + selfDir, peerDir := newPinnedPeerDirs(t) + peerMesh := clustertrust.Open(peerDir) + + const frameCount = 25 + received := make(chan []byte, frameCount) + mux := http.NewServeMux() + mux.HandleFunc(eventsPath, func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(io.LimitReader(r.Body, 1<<20)) + if err != nil { + http.Error(w, "failed to read frame", http.StatusBadRequest) + return + } + received <- body + w.WriteHeader(http.StatusOK) + }) + ts := httptest.NewUnstartedServer(mux) + ts.TLS = peerMesh.ServerTLSConfig() + ts.StartTLS() + t.Cleanup(ts.Close) + + m := newBroadcastManagerForPeer(t, selfDir, ts) + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + go m.broadcastLoop(ctx) + + for i := 0; i < frameCount; i++ { + m.broadcastFrame("workload:started", []byte(fmt.Sprintf(`{"seq":%d}`, i))) + } + + for want := 0; want < frameCount; want++ { + select { + case body := <-received: + assertFrame(t, body, want) + case <-ctx.Done(): + t.Fatalf("peer received %d of %d frames: %v", want, frameCount, ctx.Err()) + } + } +} + +// snapshotPauseHandler stops broadcastSnapshot after it copies activeLocal but +// before it queues the copied frames. This makes the remove race reproducible. +type snapshotPauseHandler struct { + paused chan struct{} + release <-chan struct{} + once sync.Once +} + +func (*snapshotPauseHandler) Enabled(context.Context, slog.Level) bool { return true } + +func (h *snapshotPauseHandler) Handle(_ context.Context, record slog.Record) error { + if record.Message == "snapshot copied before removal" { + h.once.Do(func() { close(h.paused) }) + <-h.release + } + return nil +} + +func (h *snapshotPauseHandler) WithAttrs([]slog.Attr) slog.Handler { return h } +func (h *snapshotPauseHandler) WithGroup(string) slog.Handler { return h } + +func TestManager_SnapshotCannotResurrectRemovedWorkload(t *testing.T) { + m := &Manager{ + activeLocal: make(map[workloadKey]workloadEvent), + broadcastCh: make(chan []byte, 2), + peers: newPeerSet(0), + } + workload := &Workload{ + ID: "7", Model: "m", Engine: "ollama", RunID: "r1", + State: StateRunning, OriginatedFrom: "node-a", CreatedAt: 1, + } + lifecycle, err := json.Marshal(lifecycleParams{WorkloadInfo: workload}) + if err != nil { + t.Fatalf("marshal lifecycle: %v", err) + } + m.trackActive(workloadKey{origin: "node-a", engine: "ollama", runID: "r1", id: "7"}, MethodStarted, lifecycle, StateRunning) + + paused := make(chan struct{}) + release := make(chan struct{}) + var releaseOnce sync.Once + releaseSnapshot := func() { releaseOnce.Do(func() { close(release) }) } + defer releaseSnapshot() + previousLogger := slog.Default() + slog.SetDefault(slog.New(&snapshotPauseHandler{paused: paused, release: release})) + t.Cleanup(func() { slog.SetDefault(previousLogger) }) + + snapshotDone := make(chan struct{}) + go func() { + m.broadcastSnapshot("snapshot copied before removal") + close(snapshotDone) + }() + select { + case <-paused: + case <-time.After(2 * time.Second): + t.Fatal("snapshot did not pause after copying the workload") + } + + removal, err := json.Marshal(removeParams{WorkloadID: "7", OriginatedFrom: "node-a"}) + if err != nil { + t.Fatalf("marshal removal: %v", err) + } + removeDone := make(chan struct{}) + go func() { + m.handleLocalRemove(&Message{Method: MethodRemove, Params: removal}) + close(removeDone) + }() + // The current implementation queues the removal while the snapshot is + // paused. An implementation that serializes both operations may instead + // hold the removal until the snapshot resumes. + select { + case <-removeDone: + case <-time.After(100 * time.Millisecond): + } + releaseSnapshot() + for _, done := range []<-chan struct{}{snapshotDone, removeDone} { + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("snapshot or removal did not finish") + } + } + + // The peer already has the running workload. Apply the queued frames in + // their actual delivery order through the receiver's handlers. + present := true + removes := 0 + receiver := NewServer(0, newDedupIndex(8), nil, + func(*Workload) error { present = true; return nil }, + func(string, string) error { present = false; removes++; return nil }, + nil, + ) + for len(m.broadcastCh) > 0 { + var frame Message + if err := json.Unmarshal(<-m.broadcastCh, &frame); err != nil { + t.Fatalf("decode queued frame: %v", err) + } + response := httptest.NewRecorder() + switch frame.Method { + case MethodStarted: + receiver.handleLifecycle(response, &frame) + case MethodRemove: + receiver.handleRemove(response, &frame) + default: + t.Fatalf("unexpected queued method %q", frame.Method) + } + if response.Code != http.StatusOK { + t.Fatalf("peer handled %s with HTTP %d", frame.Method, response.Code) + } + } + if removes != 1 { + t.Fatalf("peer received %d removals, want 1", removes) + } + if present { + t.Fatal("stale snapshot upsert resurrected the workload after its removal") + } +} diff --git a/services/nvpair-workload-manager/progress_test.go b/services/nvpair-workload-manager/progress_test.go index ff7f912d5..1e928d24e 100644 --- a/services/nvpair-workload-manager/progress_test.go +++ b/services/nvpair-workload-manager/progress_test.go @@ -80,7 +80,12 @@ func newProgressPeer(t *testing.T) *progressPeer { broadcaster: NewBroadcaster(peers, self), activeLocal: make(map[workloadKey]workloadEvent), ctx: context.Background(), + broadcastCh: make(chan []byte, broadcastQueueDepth), } + // Frames go out through the ordered broadcast worker, as in Run. + ctx, cancel := context.WithCancel(context.Background()) + go m.broadcastLoop(ctx) + t.Cleanup(cancel) t.Cleanup(m.broadcaster.CloseIdle) p.manager = m return p diff --git a/services/nvpair-workload-manager/server.go b/services/nvpair-workload-manager/server.go index af42a04ae..a943b26fd 100644 --- a/services/nvpair-workload-manager/server.go +++ b/services/nvpair-workload-manager/server.go @@ -179,13 +179,23 @@ func (s *Server) handleLifecycle(w http.ResponseWriter, msg *Message) { // discovery backfill), which intentionally re-asserts the same key and must // reach the broker so its store can reconcile (e.g. un-stick a wrongly // inferred failed). The store is idempotent, so bypassing dedup here is safe. - if !isResyncFrame(msg.Params) && s.dedup.seenOrAdd(keyLifecycle(wl)) { + // + // The key is recorded only after the broker emit succeeds. A failed emit + // leaves the key available for a retry; concurrent requests for the same key + // wait for that result rather than both emitting to the broker. + resync := isResyncFrame(msg.Params) + var duplicate bool + if !resync { + duplicate, err = s.dedup.emitOnce(keyLifecycle(wl), func() error { return s.emitUpsert(wl) }) + } else { + err = s.emitUpsert(wl) + } + if duplicate { slog.Debug("inter-node lifecycle deduplicated", "method", msg.Method, "id", wl.ID, "state", wl.State) s.ok(w) return } - - if err := s.emitUpsert(wl); err != nil { + if err != nil { // A failed stdout write means the broker is gone; report a server // error so the peer's retry budget can kick in, but the local // interface severing is handled as a shutdown signal elsewhere. @@ -237,13 +247,17 @@ func (s *Server) handleRemove(w http.ResponseWriter, msg *Message) { return } - if s.dedup.seenOrAdd(keyRemove(nodeID, workloadID)) { + // A failed emit leaves the key available for a retry. Concurrent requests + // for the same removal wait, so only one successful emit reaches the broker. + duplicate, err := s.dedup.emitOnce(keyRemove(nodeID, workloadID), func() error { + return s.emitRemove(workloadID, nodeID) + }) + if duplicate { slog.Debug("inter-node removal deduplicated", "workloadId", workloadID, "node", nodeID) s.ok(w) return } - - if err := s.emitRemove(workloadID, nodeID); err != nil { + if err != nil { slog.Error("failed to emit workloads:remove", "workloadId", workloadID, "node", nodeID, "err", err) http.Error(w, "broker unavailable", http.StatusInternalServerError) return diff --git a/services/nvpair-workload-manager/server_test.go b/services/nvpair-workload-manager/server_test.go new file mode 100644 index 000000000..cda57e777 --- /dev/null +++ b/services/nvpair-workload-manager/server_test.go @@ -0,0 +1,154 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "errors" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" +) + +// TestServer_InterNodeDedupRecordedOnlyAfterSuccessfulEmit verifies that a +// failed broker emit does not record a dedup key, so the peer can retry, while +// a successful emit prevents later duplicates from reaching the broker. +func TestServer_InterNodeDedupRecordedOnlyAfterSuccessfulEmit(t *testing.T) { + test := func(name string, frame []byte, wantUpserts, wantRemoves int) { + t.Run(name, func(t *testing.T) { + self, peer := newPinnedPeerMeshes(t) + var upserts, removes int + failEmit := true + srv := NewServer(0, newDedupIndex(100), self, + func(*Workload) error { + if failEmit { + return errors.New("broker gone") + } + upserts++ + return nil + }, + func(string, string) error { + if failEmit { + return errors.New("broker gone") + } + removes++ + return nil + }, + nil, + ) + post := serveEventsOverMTLS(t, srv, self, peer) + + if code := post(frame); code != http.StatusInternalServerError { + t.Fatalf("failed emit status = %d, want 500", code) + } + if upserts != 0 || removes != 0 { + t.Fatalf("emits after failed attempt = (%d upserts, %d removes), want (0, 0)", upserts, removes) + } + + failEmit = false + if code := post(frame); code != http.StatusOK { + t.Fatalf("retry status = %d, want 200", code) + } + if upserts != wantUpserts || removes != wantRemoves { + t.Fatalf("emits after retry = (%d upserts, %d removes), want (%d, %d)", upserts, removes, wantUpserts, wantRemoves) + } + // they are duplicates at this point and shouldn't get sent to do work + if code := post(frame); code != http.StatusOK { + t.Fatalf("duplicate status = %d, want 200", code) + } + if upserts != wantUpserts || removes != wantRemoves { + t.Fatalf("emits after duplicate = (%d upserts, %d removes), want (%d, %d)", upserts, removes, wantUpserts, wantRemoves) + } + }) + } + + lifecycleFrame := []byte(`{"jsonrpc":"2.0","method":"workload:started","params":` + + `{"workloadInfo":{"id":"7","model":"llama3","engine":"ollama",` + + `"runId":"r1","state":"running","originatedFrom":"uuid-peer"}}}`) + removeFrame := []byte(`{"jsonrpc":"2.0","method":"workloads:remove",` + + `"params":{"workloadId":"7","originatedFrom":"uuid-peer"}}`) + + test("lifecycle upsert", lifecycleFrame, 1, 0) + test("workload removal", removeFrame, 0, 1) +} + +// TestServer_ConcurrentDuplicatesEmitOnce verifies that an in-flight event +// reserves its key until the broker emit finishes. Both paths share the same +// dedup index, but lifecycle and removal use distinct keys and emitters. +func TestServer_ConcurrentDuplicatesEmitOnce(t *testing.T) { + test := func(name string, params []byte, handle func(*Server, http.ResponseWriter, *Message)) { + t.Run(name, func(t *testing.T) { + // Pause the first broker emit after it reserves the key. This gives + // the identical second request a chance to reach the same handler. + firstEntered := make(chan struct{}) + releaseFirst := make(chan struct{}) + // A failure before the explicit release must not strand that goroutine. + defer func() { + select { + case <-releaseFirst: + default: + close(releaseFirst) + } + }() + var emits atomic.Int32 + emit := func() error { + if emits.Add(1) == 1 { + close(firstEntered) + <-releaseFirst + } + return nil + } + srv := NewServer(0, newDedupIndex(100), nil, + func(*Workload) error { return emit() }, + func(string, string) error { return emit() }, + nil, + ) + msg := &Message{Params: params} + post := func(result chan<- int) { + recorder := httptest.NewRecorder() + handle(srv, recorder, msg) + result <- recorder.Code + } + firstDone := make(chan int, 1) + secondDone := make(chan int, 1) + go post(firstDone) + // Wait until the first request is inside emit, then start its duplicate. + select { + case <-firstEntered: + case <-time.After(2 * time.Second): + t.Fatal("first request did not reach the broker emit") + } + go post(secondDone) + // The duplicate must wait for the first emit's result. The short + // timeout gives it an opportunity to expose a premature response. + select { + case code := <-secondDone: + t.Fatalf("duplicate completed before first emit: HTTP %d", code) + case <-time.After(50 * time.Millisecond): + } + // Once the first emit succeeds, both requests may return 200, but + // only that first request should have called the broker emitter. + close(releaseFirst) + for _, result := range []<-chan int{firstDone, secondDone} { + select { + case code := <-result: + if code != http.StatusOK { + t.Fatalf("request status = %d, want 200", code) + } + case <-time.After(2 * time.Second): + t.Fatal("request did not finish") + } + } + if got := emits.Load(); got != 1 { + t.Fatalf("broker emits = %d, want 1", got) + } + }) + } + + lifecycle := []byte(`{"workloadInfo":{"id":"7","model":"llama3","engine":"ollama","runId":"r1","state":"running","originatedFrom":"uuid-peer"}}`) + remove := []byte(`{"workloadId":"7","originatedFrom":"uuid-peer"}`) + test("lifecycle upsert", lifecycle, (*Server).handleLifecycle) + test("workload removal", remove, (*Server).handleRemove) +} diff --git a/services/versions.json b/services/versions.json index f85176275..c87972651 100644 --- a/services/versions.json +++ b/services/versions.json @@ -8,14 +8,14 @@ "nvpair-node-info": "1.3.0", "nvpair-node-scanner": "1.2.1", "nvpair-manual-nodes": "1.2.2", - "nvpair-workload-manager": "1.3.0", + "nvpair-workload-manager": "1.3.1", "nvpair-errors": "1.1.1", "nvpair-node-settings": "3.1.0", "nvpair-ui-broker": "2.4.1", "nvpair-engine-manager": "1.4.0", "nvpair-cluster-manager": "2.1.1", "nvpair-job-scheduler": "1.1.0", - "nvpair-tui": "1.2.0", + "nvpair-tui": "1.2.1", "omlx-proxy": "1.4.2" } }