feat(provider): graceful drain for darkbloom stop/start/restart#392
feat(provider): graceful drain for darkbloom stop/start/restart#392Gajesh2007 wants to merge 14 commits into
Conversation
When a user runs darkbloom stop, start, or restart, drain active inference before killing or replacing the daemon. This avoids aborting in-flight requests mid-stream and improves fleet reliability. - AppKit host cancels the serve task on SIGTERM/SIGINT and waits for ProviderLoop.run() to finish its existing shutdown drain. - CLI commands read the daemon PID from DaemonStateFile, send SIGTERM, and wait up to 120s before force-killing. - User sees "Provider is currently serving requests..." when inferenceActive is true in the state file. - LaunchAgent plist now sets ExitTimeOut=600 so launchd gives the daemon the same 10-minute drain window before SIGKILL. - Add ProcessLifecycle.stopProcessGracefully and ProviderDrainHelper.drainRunningProvider helpers. Tests: ProcessLifecycleTests (PID lock + graceful stop), LaunchAgentServicePlistTests.
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
This PR refactors the graceful-shutdown and signal-handling paths in the Swift provider — the security-relevant changes are narrow and neutral-to-positive, with one new attack surface worth documenting. Trust Boundaries Touched
Per-Threat AssessmentT-008 — Provider sends plaintext SSE chunks on encryption failureℹ️ Neutral. The shutdown drain pathway ( T-009 — Swift provider excluded from private-request routing (Python flags)ℹ️ Neutral. No changes to the privacy-capability flags sent during registration or the routing gate logic. SEC-017 unchanged. T-010 — In-flight inference cancellation not propagated to inference engine✅ Strengthens mitigation. The new T-034 — Provider runs modified code while advertising a trusted identityℹ️ Neutral. The T-042 — Forged code-identity via attacker-key challenge + log/disk harvestℹ️ Neutral, with one structural observation. The new T-015 — Operator disables SIP to bypass security controlsℹ️ Neutral. Signal-source installation suppresses the OS default action for SIGTERM/SIGINT ( New Attack Surface Not Covered by an Existing ThreatRace between In Suggested fix: Set source.setEventHandler { [weak self] in
self?.terminationRequested = true // set before cancel/terminate
self?.serveTask?.cancel()
NSApp.terminate(nil)
}Open FindingsNo SEC-* findings are resolved by this PR. The changes are infrastructure (shutdown sequencing) and do not address SEC-016, SEC-017, SEC-007, SEC-035, or any other open finding. 🔐 Threat model: |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 263bbe87eb
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| /// `reply(toApplicationShouldTerminate: true)` once the drain finishes. | ||
| func applicationShouldTerminate(_ sender: NSApplication) -> NSApplication.TerminateReply { | ||
| terminationRequested = true | ||
| serveTask?.cancel() |
There was a problem hiding this comment.
When darkbloom stop sends SIGTERM while inference is active, this cancellation propagates into ProviderLoop.run(). The shutdown path then calls waitForInflightDrain, but that helper returns false immediately when Task.isCancelled is set and proceeds to cancelAllInflight(), so the new graceful-stop path still aborts active streams instead of waiting for them to finish. The stop/terminate path needs to trigger coordinator shutdown without leaving the provider loop task in the cancelled state used by the drain loop.
Useful? React with 👍 / 👎.
| exit(0) | ||
| await self.finishServeTask() | ||
| } catch { | ||
| await self.finishServeTask() |
There was a problem hiding this comment.
For hosted start --foreground / --local, any startup or runtime error now enters this catch, but finishServeTask() calls exit(0) whenever termination was not requested, so Darkbloom.exit(withError:) is never reached. In failure cases such as bad arguments, GPU/model initialization errors, or a provider loop error, the daemon exits successfully and can hide the actual failure from launchd/scripts. Handle the error path separately or make finishServeTask avoid exiting successfully from catch.
Useful? React with 👍 / 👎.
| try? await Task.sleep(nanoseconds: 50_000_000) // 0.05s | ||
| } | ||
|
|
||
| return processIsAlive(pid) ? .forceKilled(pid) : .stoppedGracefully |
There was a problem hiding this comment.
After the graceful timeout expires, this path sends SIGKILL but returns .stoppedGracefully as soon as that kill succeeds. drainRunningProvider only warns on .forceKilled, and stop can then say the provider stopped after draining, so requests that exceeded the 120s drain window are reported as cleanly drained even though they were killed. Return a force-kill outcome whenever the timeout escalation path is taken.
Useful? React with 👍 / 👎.
| action: DrainAction, | ||
| timeout: TimeInterval = 120.0 | ||
| ) async -> Bool { | ||
| guard LaunchAgent.isLoaded() else { return false } |
There was a problem hiding this comment.
On upgraded machines still loaded under the legacy dev.darkbloom.provider label, this guard returns false even though LaunchAgent.stop()/restart() still operate on legacy labels, so those commands fall through to the launchd bootout/kickstart path without waiting on the daemon PID first. Active requests on legacy-labeled providers therefore bypass the new graceful-drain path; use the supported-label check here and in the parallel start replacement check before manipulating launchd.
Useful? React with 👍 / 👎.
| public static func processIsAlive(_ pid: Int32) -> Bool { | ||
| let rc = kill(pid, 0) | ||
| if rc == 0 { return true } | ||
| return errno != ESRCH |
There was a problem hiding this comment.
This helper is now fed by the JSON daemon state file before sending real signals, but it does not reject pid <= 0. If that state file is malformed or corrupted while still fresh, pid == 0 makes the new stop path signal the caller's process group and pid == -1 can signal every process the user is allowed to signal, rather than just the provider daemon. Add the same positive-PID guard used by daemonProcessAlive before any kill call.
Useful? React with 👍 / 👎.
Add inflightRequestCount to DaemonStateFile and populate it from ProviderLoop so the CLI can tell the user exactly how many requests are in flight and how long it will wait before force-killing.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8c66f6399b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if let state, !state.isStale(now: now), ProcessLifecycle.processIsAlive(state.pid) { | ||
| pid = state.pid |
There was a problem hiding this comment.
When the launchd job is still loaded but the daemon exited or crashed within the 90s state freshness window, this trusts the old state-file PID solely because kill(pid, 0) succeeds. If macOS has already reused that PID, stop/restart will send SIGTERM/SIGKILL to an unrelated process instead of falling back to launchd. Verify the PID against launchd’s reported process, the provider PID file, or the executable identity before signalling it.
Useful? React with 👍 / 👎.
| public var currentModel: String? | ||
| public var warmModels: [String] | ||
| public var inferenceActive: Bool | ||
| public var inflightRequestCount: Int |
There was a problem hiding this comment.
Adding this non-optional Codable field while keeping schema 1 makes DaemonStateFile.read() fail with keyNotFound for schema-1 state files written by existing providers, because they do not contain inflight_request_count. Until the daemon rewrites the file, darkbloom status/doctor can report a live daemon as not running and the new drain helper cannot read its PID. Decode the field with a default, or migrate/bump the schema so existing state files remain readable.
Useful? React with 👍 / 👎.
- P1: waitForInflightDrain no longer aborts on task cancellation once the controlled shutdown path has started, so SIGTERM/SIGINT-triggered drains actually wait for active inference instead of falling through to cancelAllInflight immediately. - P2: hosted serve failures no longer mask nonzero exits; finishServeTask is only called when AppKit requested termination. - P2: stopProcessGracefully now returns .forceKilled whenever the drain timeout expires and SIGKILL is used. - P2: ProviderDrainHelper and StartCommand use isAnySupportedLabelLoaded so legacy-labeled providers are drained too. - P2: stopProcessGracefully rejects pid <= 0 before calling kill. - Added tests for nonpositive PID guard and SIGTERM-ignored force kill.
Add a hasReplied flag so AppKit's reply(toApplicationShouldTerminate:) is only called once if termination is requested multiple times.
|
Found 1 test failure on Blacksmith runners: Failure
|
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 423111c4da
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| func drainRunningProvider( | ||
| action: DrainAction, | ||
| timeout: TimeInterval = 120.0 | ||
| ) async -> Bool { |
There was a problem hiding this comment.
When an in-flight stream takes longer than two minutes, this default makes every darkbloom stop / restart / replacement start escalate to SIGKILL at 120s via stopProcessGracefully, even though ProviderLoop now waits up to 600s to drain and the launchd plist was updated with ExitTimeOut = 600. In that scenario the CLI path aborts requests that the daemon would have completed cleanly; use the same drain window or reserve force-kill behavior for an explicit force path.
Useful? React with 👍 / 👎.
| func applicationShouldTerminate(_ sender: NSApplication) -> NSApplication.TerminateReply { | ||
| terminationRequested = true | ||
| serveTask?.cancel() | ||
| return .terminateLater |
There was a problem hiding this comment.
When SIGTERM arrives during a coordinator-routed stream, cancelling the serve task here drives ProviderLoop's cancellation handler, which calls CoordinatorClient.shutdown() before waitForInflightDrain; that shutdown closes the WebSocket and finishes the outbound router, so chunks and the final completion produced while the daemon is “draining” are dropped instead of reaching the consumer. Fresh evidence in this revision is that the wait loop now continues after cancellation, but the transport is still torn down first, so the graceful stop path still aborts active coordinator streams.
Useful? React with 👍 / 👎.
| let outcome = await ProcessLifecycle.stopProcessGracefully(pid: pid, timeout: 2.0) | ||
|
|
||
| #expect(outcome == .stoppedGracefully) |
There was a problem hiding this comment.
When this test targets a Process child, SIGTERM can make /bin/sleep exit but leave it as a zombie until the parent calls waitUntilExit; processIsAlive is based on kill(pid, 0), which still succeeds for zombies, so stopProcessGracefully waits the full timeout and returns .forceKilled instead of the expected graceful outcome. Reap the spawned process (or test against a non-child helper) before using kill(pid, 0) as the disappearance signal.
Useful? React with 👍 / 👎.
- Increase CLI/daemon graceful-stop timeout default to 600s to match launchd ExitTimeOut. - Cross-check daemon state-file PID against ~/.darkbloom/provider.pid before signalling. - Reap spawned child processes in ProcessLifecycle graceful-stop tests. - Fix ProviderLoop shutdown to use coordinatorClient? and keep socket open during drain.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e2f3673934
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // Cancel background work first, but keep the coordinator socket open | ||
| // so in-flight responses can still be delivered while we drain. | ||
| await cancelBackgroundWorkAndPreloads() | ||
| let drained = await waitForInflightDrain(timeout: Self.shutdownDrainTimeout) |
There was a problem hiding this comment.
When SIGTERM cancels the run task, the for await event in events loop has ended, but this wait keeps CoordinatorClient.shutdown() deferred until the in-flight drain completes. During a long stop/restart drain, CoordinatorClient.receiveLoop and heartbeats can therefore keep the WebSocket registered and yield new valid inference_request events into an unconsumed stream, so those requests never reach handleInferenceRequest to get the new 503 shutdown response and will wait for coordinator timeout instead of rerouting. Keep a consumer alive to reject/cancel inbound work during the drain, or mark the provider unavailable before leaving the socket open.
Useful? React with 👍 / 👎.
Address Codex review: during a graceful stop/restart drain the run loop has
stopped consuming coordinator events, but the WebSocket is intentionally kept
open so in-flight responses can finish. New inference_requests routed by the
coordinator during a long drain were yielded into the unconsumed event stream
and never answered, so the consumer waited for an inference timeout instead of
being rerouted.
CoordinatorClient.beginDraining() now flips a flag the receive loop checks: new
inference_requests are answered with a 503 ("provider is shutting down") right
in the receive loop (the socket and its receive task outlive run-task
cancellation), so the coordinator reroutes immediately. ProviderLoop calls it
at the start of the drain, before waiting for in-flight work.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1edaf69831
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
|
||
| // Drain (if any) is complete: now it is safe to close the transport and | ||
| // release GPU/model resources. | ||
| await coordinatorClient?.shutdown() |
There was a problem hiding this comment.
When darkbloom stop/restart drains a stream, waitForInflightDrain only waits for finishInflightRequest to remove the request, but the inference task's send.send(...) calls merely enqueue chunks/final inference_complete into OutboundRouter's AsyncStream. On a slow WebSocket or for the final completion message, those queued frames may not have been written yet; calling shutdown() here sets shutdownRequested and cancels the WebSocket, so the outbound writer can break before sending buffered responses and the consumer/billing path may miss the tail of a request that supposedly drained successfully. The drain needs to wait for outbound delivery/flush before tearing down the coordinator transport.
Useful? React with 👍 / 👎.
| public func beginDraining() { | ||
| draining = true | ||
| } |
There was a problem hiding this comment.
Once this drain mode is enabled, ProviderLoop has stopped consuming the event stream, but CoordinatorClient still handles coordinator cancel messages by yielding .cancel into that now-unconsumed stream. If a consumer aborts an active stream during a long stop/restart drain, handleCancellation never runs, so generation continues until it finishes or hits the 600s drain timeout, delaying shutdown and wasting GPU. Draining mode should process cancels for existing in-flight requests directly, or keep a consumer alive for cancellation events.
Useful? React with 👍 / 👎.
| let loadSelf = self | ||
| do { | ||
| try await ensureModelLoaded(modelId: modelId) | ||
| try await Task.detached { try await loadSelf.ensureModelLoaded(modelId: modelId) }.value |
There was a problem hiding this comment.
When SIGTERM arrives after a coordinator request has been accepted for a cold model and the model load stalls past the 600s drain window, this detached load is not a child of the cancelled serve task and no handle is retained to cancel it. The timeout path removes requestToModel and cancels waiters, but handleInferenceRequest is still awaiting .value here (the new comment notes that a cancelled run task keeps waiting), so ProviderLoop.run() cannot return and AppKit never replies; the CLI then escalates to SIGKILL instead of exiting after the drain timeout. Retain/cancel the load task on timeout or make the detached load observe a cancellation token.
Useful? React with 👍 / 👎.
ethenotethan
left a comment
ethenotethan
left a comment
There was a problem hiding this comment.
Automated Code Review — Layr-Labs/d-inference#
Verdict: REQUEST_CHANGES
Security — ✅ No issues found
Performance — 4 finding(s) (1 blocking)
- 🟡 [MEDIUM]
provider-swift/Sources/ProviderCore/ProviderLoop.swift:1200-1217— Detached task for model loading lacks timeout or cancellation mechanism- Suggestion: Add timeout or cancellation mechanism to the detached task to prevent resource leaks if loads hang indefinitely
- 🔵 [INFO]
provider-swift/Sources/ProviderCore/ProviderLoop.swift:3308-3330— Sleep in drain polling loop blocks actor thread- Suggestion: Use async timer or yield instead of Task.sleep to avoid blocking the actor thread during drain polling
- 🔵 [INFO]
provider-swift/Sources/darkbloom/ProviderDrainHelper.swift:69— Synchronous stopProcessGracefully call blocks caller thread- Suggestion: Consider making this function async or using a background queue for the process monitoring to avoid blocking the main thread
- 🔵 [INFO]
provider-swift/Sources/ProviderCore/Service/LaunchAgent.swift:90-120— Synchronous process execution in servicePID method- Suggestion: Make servicePID async and use async process execution to avoid blocking the caller thread
Type_diligence — ✅ No issues found
Additive_complexity — 5 finding(s) (1 blocking)
- 🔵 [INFO]
provider-swift/Sources/ProviderCore/ProviderLoop.swift:222— Memoized shutdown task adds complexity for single-use scenario- Suggestion: Consider using a simple boolean flag instead of memoizing the Task since shutdown only happens once per process lifecycle
- 🟡 [MEDIUM]
provider-swift/Sources/ProviderCore/ProviderLoop.swift:1200-1216— Detached task wrapper adds cognitive load without clear necessity- Suggestion: Document why detachment is required or consider if the same effect could be achieved with simpler cancellation handling
- 🔵 [INFO]
provider-swift/Sources/ProviderCore/ProviderLoop.swift:2217-2233— Complex conditional logic in cancelLoadWaiters duplicates sparing pattern- Suggestion: Extract the sparing logic into a helper method to reduce cognitive load and improve testability
- 🔵 [INFO]
provider-swift/Sources/darkbloom/ProviderDrainHelper.swift:1-81— Entire helper file for what could be inline logic- Suggestion: Consider inlining this logic into the command files since it's only used in 3 places and adds another abstraction layer
- 🔵 [INFO]
provider-swift/Sources/darkbloom/ProviderAppKitHost.swift:54-60— Multiple boolean flags tracking termination state- Suggestion: Combine terminationRequested and hasReplied into a single enum state to reduce state management complexity
9 finding(s) total, 2 blocking. Verdict: REQUEST_CHANGES.
🤖 Automated review by Centaur · DAR-186
Address the latest Codex review round + the blocking review finding: - (BLOCKING) Cancel detached cold-load tasks on teardown. The per-request detached `ensureModelLoaded` is now retained in `inflightModelLoadTasks` and cancelled by `cancelAllInflight`. Without a handle, a load that stalled past the 600s drain timeout left `handleInferenceRequest` awaiting `loadTask.value` forever, so `run()` never returned and the CLI escalated to SIGKILL. - Handle coordinator `cancel` frames while draining. `beginDraining(onCancel:)` routes cancels straight to `handleCancellation` (the provider event stream is no longer consumed during the drain), so an aborted in-flight request stops generating promptly instead of running to the drain timeout and wasting GPU. - Flush queued outbound frames before closing the socket. `OutboundRouter` now tracks pending writes (marked written by the outbound task); the drain calls `CoordinatorClient.flushOutbound(timeout:)` before `shutdown()` so the tail of a finished request (last chunks + `inference_complete`) is delivered instead of being dropped when the WebSocket is cancelled.
ethenotethan
left a comment
ethenotethan
left a comment
There was a problem hiding this comment.
Automated Code Review — Layr-Labs/d-inference#
Verdict: REQUEST_CHANGES
Security — ✅ No issues found
Performance — 4 finding(s) (3 blocking)
- 🔵 [INFO]
provider-swift/Sources/ProviderCore/Coordinator/CoordinatorClient.swift:574-588— Unbounded pendingWrites counter without overflow protection- Suggestion: Add overflow protection:
pendingWrites = min(pendingWrites + 1, Int.max - 1000)or use saturating arithmetic to prevent integer overflow in high-throughput scenarios
- Suggestion: Add overflow protection:
- 🟡 [MEDIUM]
provider-swift/Sources/ProviderCore/Coordinator/CoordinatorClient.swift:695-703— Blocking sleep in flushOutbound with tight polling loop- Suggestion: Use exponential backoff or longer sleep intervals: start with 20ms but increase to 100-200ms after several iterations to reduce CPU usage during long flushes
- 🟡 [MEDIUM]
provider-swift/Sources/ProviderCore/ProviderLoop.swift:3345-3357— Tight polling loop in waitForInflightDrain with 250ms sleep- Suggestion: Use exponential backoff: start with 250ms but increase to 1-2 seconds after initial polls to reduce CPU overhead during long drains
- 🟡 [MEDIUM]
provider-swift/Sources/darkbloom/ProviderDrainHelper.swift:61-69— Tight polling loop in stopProcessGracefully with 100ms sleep- Suggestion: Use exponential backoff: start with 100ms but increase sleep duration (e.g., min(sleep * 1.5, 1000ms)) to reduce CPU usage during long drain waits
Type_diligence — ✅ No issues found
Additive_complexity — 5 finding(s) (2 blocking)
- 🔵 [INFO]
provider-swift/Sources/ProviderCore/Coordinator/CoordinatorClient.swift:283-289— pendingWrites counter adds complexity for minimal benefit- Suggestion: Consider if the flush timeout is sufficient without precise write tracking - the 5s timeout may be adequate for most cases
- 🟡 [MEDIUM]
provider-swift/Sources/ProviderCore/ProviderLoop.swift:258-265— inflightModelLoadTasks duplicates existing task tracking pattern- Suggestion: Unify with inflightTasks using a common task registry or extend inflightTasks to handle both inference and load tasks
- 🟡 [MEDIUM]
provider-swift/Sources/ProviderCore/ProviderLoop.swift:1220-1240— detached task pattern adds complexity to avoid cancellation inheritance- Suggestion: Consider if the graceful drain could work with regular tasks by checking shutdown state instead of detaching
- 🔵 [INFO]
provider-swift/Sources/ProviderCore/ProviderLoop.swift:2238-2258— cancelLoadWaiters has complex conditional logic that could be simplified- Suggestion: Extract the sparing logic into a helper method or use a more direct approach to distinguish background vs inflight loads
- 🔵 [INFO]
provider-swift/Sources/darkbloom/ProviderDrainHelper.swift:33-81— complex state freshness and PID validation logic for user messages- Suggestion: Simplify to just check if requests are active without complex state validation - the drain will work regardless of message accuracy
9 finding(s) total, 5 blocking. Verdict: REQUEST_CHANGES.
🤖 Automated review by Centaur · DAR-186
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b45b531fd2
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| await coordinatorClient?.beginDraining(onCancel: { requestId in | ||
| await drainSelf.handleCancellation(requestId: requestId, receivedFromCoordinator: true) | ||
| }) |
There was a problem hiding this comment.
When a cancel frame arrives during a graceful stop/restart for a request that was accepted but is still awaiting the detached ensureModelLoaded task, this new drain handler calls handleCancellation, which removes requestToModel but does not cancel inflightModelLoadTasks[requestId]. The drain can then observe no in-flight work and proceed to close/unload while the original handleInferenceRequest is still awaiting loadTask.value; if that cold load stalls, the serve task does not return and the CLI still escalates to SIGKILL. Cancel the matching inflight model-load task here or in handleCancellation.
Useful? React with 👍 / 👎.
| for candidate in supportedLabels { | ||
| if let pid = servicePID(label: candidate) { | ||
| return pid | ||
| } |
There was a problem hiding this comment.
When both the canonical and legacy launchd jobs are loaded during an upgrade, this helper returns the first running PID only. The stop/replacement paths later unload every supported label, so an active request on the second loaded label never gets the SIGTERM plus 600s drain wait and is instead torn down by the launchd bootout path. Collect and drain all supported-label PIDs before manipulating launchd.
Useful? React with 👍 / 👎.
Follow-up to the latest Codex review: - handleCancellation now also cancels the request's detached cold-load task (inflightModelLoadTasks[requestId]). A cancel frame arriving during the drain for a request still awaiting `loadTask.value` would otherwise leave the serve task blocked if the load stalled, so run() never returns and the CLI escalates to SIGKILL. Centralizing it here covers both the coordinator-cancel and drain-cancel paths (cancelAllInflight still bulk-cancels cold loads too). - drainRunningProvider now drains every loaded supported-label PID, not just the first. During an upgrade both the canonical and a legacy launchd job can be loaded; the stop/replace paths unload every label, so an active request on the second label would otherwise skip the graceful SIGTERM+drain and be torn down by bootout. LaunchAgent.loadedServicePIDs() returns all distinct running PIDs; they are drained concurrently so total wait stays bounded by the timeout.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 018e0a2477
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| @@ -775,33 +895,14 @@ public actor ProviderLoop { | |||
| } | |||
| let preloads = Array(preloadTasks.values) | |||
| for task in preloads { task.cancel() } | |||
There was a problem hiding this comment.
When a coordinator load_model is already warming model M, an inference request for M can be accepted and wait in loadingWaiters while the preload task owns the actual ensureModelLoaded. During a graceful stop/restart, sparingInflightLoads only spares the waiter, but this line still cancels every preload task, so the shared load throws and the accepted request returns a load failure instead of draining. Fresh evidence in this revision is that preloads are still canceled unconditionally even when requestToModel contains the same model; skip or transfer those preloads before canceling background-only work.
Useful? React with 👍 / 👎.
| if let loadTask = inflightModelLoadTasks.removeValue(forKey: requestId) { | ||
| loadTask.cancel() |
There was a problem hiding this comment.
When a coordinator cancel arrives while an accepted request is still cold-loading, this now cancels the detached load task, but handleInferenceRequest catches that CancellationError and sends inference_error with error: "model load failed: ..." and status 500. The coordinator's handleInferenceError treats that string as a dispatch load failure/cooldown, so a normal client abort or drain-time cancel during cold load can temporarily deroute a healthy provider/model pair; suppress the terminal or map it to a cancel/503 path instead.
Useful? React with 👍 / 👎.
| await coordinatorClient?.beginDraining(onCancel: { requestId in | ||
| await drainSelf.handleCancellation(requestId: requestId, receivedFromCoordinator: true) | ||
| }) |
There was a problem hiding this comment.
During a long graceful stop/restart, this installs only a cancel handler after the event loop has stopped consuming coordinator events. If the WebSocket disconnects in that window, CoordinatorClient still yields .disconnected into the unconsumed stream and reconnects while the provider keeps generating for a request the coordinator already failed on disconnect; the final frames then target an unknown request and the CLI waits for work that can no longer reach the consumer. Add a drain-time disconnect path that cancels in-flight work or prevents reconnecting while shutting down.
Useful? React with 👍 / 👎.
Address the latest Codex review round (all follow-ons of the drain work): - Don't cancel preloads backing accepted requests. During a drain, a preload can be the actual loader for a model an accepted request is waiting on (request parked in loadingWaiters). cancelBackgroundWorkAndPreloads now spares preloads whose model is in requestToModel and only cancels/awaits the background-only ones, so the accepted request drains instead of failing. - Treat cold-load cancels as cancellations, not faults. A cancel (client abort or drain) that cancels the detached load surfaced as inference_error "model load failed" / 5xx, which the coordinator treats as a model load failure and cools down a healthy provider+model. The catch now maps CancellationError to a 503 reroute while shutting down, and stays silent for a plain cancel (the coordinator already aborted it). - Cancel in-flight work if the socket drops mid-drain. beginDraining(onDisconnect:) lets CoordinatorClient.runLoop, on a disconnect while draining, cancel in-flight work (the coordinator already failed those requests) and stop reconnecting, instead of generating frames that can never reach the consumer.
ethenotethan
left a comment
ethenotethan
left a comment
There was a problem hiding this comment.
Automated Code Review — Layr-Labs/d-inference#
Verdict: REQUEST_CHANGES
Security — ✅ No issues found
Performance — 4 finding(s) (3 blocking)
- 🔵 [INFO]
provider-swift/Sources/ProviderCore/Coordinator/CoordinatorClient.swift:574-593— pendingWrites counter incremented without bounds checking- Suggestion: Add bounds checking or circuit breaker to prevent unbounded growth of pendingWrites counter during network congestion
- 🟡 [MEDIUM]
provider-swift/Sources/ProviderCore/ProviderLoop.swift:3379-3398— Busy polling with 250ms sleep in waitForInflightDrain- Suggestion: Replace polling with event-driven approach using async streams or condition variables to avoid CPU waste during long drains
- 🟡 [MEDIUM]
provider-swift/Sources/darkbloom/ProviderDrainHelper.swift:73-81— Concurrent task group without size limits for PID draining- Suggestion: Add concurrency limit to TaskGroup or validate PID array size before spawning tasks to prevent resource exhaustion with many loaded services
- 🟡 [MEDIUM]
provider-swift/Sources/ProviderCore/Service/ProcessLifecycle.swift:165-185— Busy polling with 100ms sleep in stopProcessGracefully- Suggestion: Use process monitoring APIs or signal handlers instead of polling to detect process termination more efficiently
Type_diligence — 2 finding(s)
- 🔵 [INFO]
provider-swift/Sources/ProviderCore/Coordinator/CoordinatorClient.swift:283— pendingWrites initialized to 0 implicitly relies on zero value- Suggestion: Add explicit initialization: private var pendingWrites = 0
- 🔵 [INFO]
provider-swift/Sources/ProviderCore/Service/DaemonStateFile.swift:117— inflightRequestCount uses optional Int? where non-optional might be clearer- Suggestion: Consider using non-optional Int with default 0 if the field should always have a value
Additive_complexity — 5 finding(s) (2 blocking)
- 🔵 [INFO]
provider-swift/Sources/ProviderCore/Coordinator/CoordinatorClient.swift:283-289— pendingWrites counter adds complexity for minimal benefit- Suggestion: Consider if the 5-second flush timeout justifies this tracking complexity, or use a simpler approach like a brief fixed delay
- 🟡 [MEDIUM]
provider-swift/Sources/ProviderCore/ProviderLoop.swift:260-264— inflightModelLoadTasks duplicates existing cancellation patterns- Suggestion: Unify with existing inflightTasks pattern or use a single map with tagged task types
- 🟡 [MEDIUM]
provider-swift/Sources/ProviderCore/ProviderLoop.swift:1247-1260— detached task pattern adds complexity to avoid cancellation inheritance- Suggestion: Consider if the graceful drain benefit justifies this complexity vs simpler approaches like checking shutdown state in load path
- 🔵 [INFO]
provider-swift/Sources/ProviderCore/ProviderLoop.swift:2753-2757— shutdownShouldAbortLoad duplicates isShuttingDown checks throughout load path- Suggestion: Consolidate shutdown checks or make the load path consistently use this helper
- 🔵 [INFO]
provider-swift/Sources/darkbloom/ProviderDrainHelper.swift:1-90— entire file could be methods on existing types rather than standalone functions- Suggestion: Move drainRunningProvider to LaunchAgent or ProcessLifecycle as static methods to keep related functionality together
11 finding(s) total, 5 blocking. Verdict: REQUEST_CHANGES.
🤖 Automated review by Centaur · DAR-186
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 85ed098ce6
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| error: "provider is shutting down", | ||
| statusCode: 503 | ||
| ) | ||
| try? await ws.send(.string(errorResponse)) |
There was a problem hiding this comment.
When a new inference_request arrives during the window after the last in-flight request finishes but before shutdown, this direct ws.send is not counted by OutboundRouter.pendingCount(). Because the actor can run flushOutbound()/shutdown() while this send is suspended, the flush can see zero pending messages and cancel the WebSocket before the 503 reject is written, leaving that coordinator request to wait for timeout instead of rerouting; route this response through the outbound router or include these drain-time sends in the flush accounting.
Useful? React with 👍 / 👎.
| // and to route `cancel` frames for in-flight requests straight to | ||
| // `handleCancellation` so an aborted stream stops generating promptly. | ||
| let drainSelf = self | ||
| await coordinatorClient?.beginDraining( |
There was a problem hiding this comment.
This only flips the coordinator client into drain mode for future WebSocket reads after the provider event loop has already stopped consuming events. If another inference_request was yielded into the AsyncStream buffer while the loop was busy with a prior cold load/generation, SIGTERM makes the loop exit before that buffered request is processed, and it is neither registered in requestToModel nor rejected with 503, so the coordinator waits for timeout instead of rerouting; drain or reject the buffered events before waiting for in-flight work.
Useful? React with 👍 / 👎.
| onCancel: { requestId in | ||
| await drainSelf.handleCancellation(requestId: requestId, receivedFromCoordinator: true) |
There was a problem hiding this comment.
When a cancel frame arrives during drain for a stream that already emitted output, this handler calls handleCancellation, which immediately removes the request from requestToModel/inflightTasks; waitForInflightDrain can then observe no work and close the socket before the canceled task emits its partial inference_complete/499 terminal. The coordinator has a settlement grace path for terminals after client-gone, so dropping that terminal refunds the reservation and loses provider payout for work already streamed; keep the request counted until the task's defer finishes or flush after the cancellation terminal is queued.
Useful? React with 👍 / 👎.
| let inflightModels = sparingInflightLoads ? Set(requestToModel.values) : Set<String>() | ||
| let preloadsToCancel = preloadTasks.filter { !inflightModels.contains($0.key) } | ||
| for (_, task) in preloadsToCancel { task.cancel() } | ||
| cancelLoadWaiters(sparingInflightRequests: sparingInflightLoads) |
There was a problem hiding this comment.
When shutdown starts while an accepted request is waiting behind an existing load_model preload, this filter spares that preload for the drain, but the timeout path only cancels inflightModelLoadTasks and waiters. If that preload is the operation that stalls past the 600s drain window, the old preloadTasks[model] is never cancelled or awaited after its backing request is force-cancelled, so a scheduled-window shutdown can return with the old ProviderLoop still loading a model in the background and racing the next window; cancel the spared preload tasks once the drain timeout fires.
Useful? React with 👍 / 👎.
2 participants
When a user runs
darkbloom stop,start, orrestart, drain active inference before killing or replacing the daemon. This avoids aborting in-flight requests mid-stream and improves fleet reliability.Before / After
%%{init: {'theme': 'base', 'themeVariables': {'primaryColor': '#ffcccc', 'primaryTextColor': '#333', 'primaryBorderColor': '#999', 'lineColor': '#666', 'secondaryColor': '#ccffcc', 'tertiaryColor': '#eeeeee'}}}%% sequenceDiagram autonumber participant U as User participant CLI as darkbloom CLI participant L as launchd participant D as Provider daemon participant C as Coordinator rect rgb(255, 220, 220) Note over U,C: BEFORE: immediate kill/restart U->>CLI: darkbloom stop / restart / start CLI->>L: bootout / kickstart -k L->>D: SIGTERM → SIGKILL (20s default) D--xC: in-flight requests aborted end rect rgb(220, 255, 220) Note over U,C: AFTER: graceful drain U->>CLI: darkbloom stop / restart / start CLI->>D: SIGTERM via daemon PID D->>D: reject new work, wait for<br/>in-flight inference to finish D->>CLI: process exits cleanly CLI->>L: bootout / kickstart endWhat changed
ProviderLoop.run()to finish its existing 10-minute shutdown drain.StopCommand,RestartCommand, andStartCommandread the daemon PID fromDaemonStateFile, sendSIGTERM, and wait up to 120s before force-killing.Provider is currently serving 3 requests. Waiting up to 120s for them to finish before stopping...ExitTimeOut: 600solaunchctl bootout/ logout shutdown gives the daemon the same 10-minute drain window before SIGKILL.ProcessLifecycle.stopProcessGracefullyandProviderDrainHelper.drainRunningProvider.Tests
ProcessLifecycleTestscovers PID lock behavior and the new graceful stop helper (including stopping a real subprocess).LaunchAgentServicePlistTestsasserts the newExitTimeOut = 600.swift buildpasses.swift testsuite did not complete in this environment due to a missing MLX metallib (pre-existing, unrelated).