Add cancellation and deadline support to Session.process_request - #993
Closed
bmehta001 wants to merge 1 commit into
Closed
Add cancellation and deadline support to Session.process_request#993bmehta001 wants to merge 1 commit into
bmehta001 wants to merge 1 commit into
Conversation
Session.process_request() previously had no way to time out or be cancelled outside the streaming path, so a non-terminating non-streaming generation permanently pinned the session refcount. This caused Model.Unload() to fail with '1 session(s) still using it', FoundryLocalManager.close() to exceed its drain deadline, and process teardown to intermittently crash. C++ core: - Add ICancellable interface and OgaGeneratorCancellable adapter so Session can interrupt ORT GenAI mid-compute (SetRuntimeOption(terminate_session)), not just between token boundaries. - Add Request::SetTimeout/ArmDeadline/ShouldStop for a re-armable wall-clock deadline that applies to streaming and non-streaming requests alike. - Add Session::Cancel() with active-generator tracking and a deadline watchdog thread; wire it into all chat/audio/embeddings generation loops. - Cancel() now latches request.canceled on every in-flight request (not just published generators) so finish_reason and history rollback are correct even when cancellation lands during prefill. - Add LiveSessionRegistry so SessionManager::CancelAll() can reach every live session, including ones created via the direct API that never took a SessionRegistration. - Add FOUNDRY_LOCAL_ERROR_TIMEOUT and two C ABI vtable entries: Request_SetTimeoutMs, Session_Cancel (appended to preserve ABI ordering). Bindings (C++ wrapper, Python, C#, JS): - Request.SetTimeout()/set_timeout()/setTimeout(), Session.Cancel()/cancel(). - Non-streaming ProcessRequestAsync (C#) and processRequest (JS) now accept a CancellationToken / AbortSignal that genuinely interrupts an in-flight generation, not just prevents scheduling. - Session teardown (Python _close, C# Dispose, JS dispose) now cancels the session itself, covering the non-streaming case that was previously invisible to shutdown. Tests: added ChatSessionTest coverage for timeout enforcement, timeout error code, deadline re-arming across reused requests, mid-flight cross-thread cancellation, cancel-before-request rejection, and cancel-when-idle safety. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 86c9e761-c751-48bc-8e73-2db5257c9e88
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Contributor
There was a problem hiding this comment.
Pull request overview
Adds cross-language cancellation and deadline support for inference sessions, including mid-compute interruption and shutdown handling.
Changes:
- Adds native request deadlines, session cancellation, watchdogs, and live-session tracking.
- Exposes timeout and cancellation APIs through C++, Python, C#, and JavaScript.
- Adds C++ regression coverage for cancellation and deadlines.
Reviewed changes
Copilot reviewed 37 out of 37 changed files in this pull request and generated 14 comments.
Show a summary per file
| File | Description |
|---|---|
sdk_v2/python/src/foundry_local_sdk/session.py |
Adds timeout, cancellation, and teardown handling. |
sdk_v2/python/src/foundry_local_sdk/request.py |
Exposes request timeouts. |
sdk_v2/python/src/foundry_local_sdk/_native/build_cffi.py |
Extends Python C ABI definitions. |
sdk_v2/js/src/session.ts |
Adds AbortSignal cancellation and session cancellation. |
sdk_v2/js/src/request.ts |
Exposes request timeouts. |
sdk_v2/js/src/detail/native.ts |
Extends native TypeScript interfaces. |
sdk_v2/js/src/detail/errors.ts |
Adds the timeout error code. |
sdk_v2/js/native/src/session.h |
Declares native session cancellation methods. |
sdk_v2/js/native/src/session.cc |
Implements native session cancellation bindings. |
sdk_v2/js/native/src/request.h |
Declares native timeout support. |
sdk_v2/js/native/src/request.cc |
Implements native timeout conversion. |
sdk_v2/cs/src/Session.cs |
Adds cancellation-token and teardown cancellation. |
sdk_v2/cs/src/Request.cs |
Exposes request deadlines. |
sdk_v2/cs/src/Detail/NativeMethods.cs |
Extends C# ABI declarations. |
sdk_v2/cs/src/Detail/FoundryLocalApi.cs |
Maps timeout errors and session cancellation. |
sdk_v2/cpp/test/internal_api/chat/chat_session_test.cc |
Adds cancellation and timeout tests. |
sdk_v2/cpp/src/inferencing/session/session.h |
Defines cancellation state and generator tracking. |
sdk_v2/cpp/src/inferencing/session/session.cc |
Implements cancellation and deadline watchdogs. |
sdk_v2/cpp/src/inferencing/session/session_manager.h |
Documents global cancellation behavior. |
sdk_v2/cpp/src/inferencing/session/session_manager.cc |
Cancels sessions during shutdown. |
sdk_v2/cpp/src/inferencing/session/request.h |
Adds deadline and stop state. |
sdk_v2/cpp/src/inferencing/session/oga_generator_cancellable.h |
Declares the OGA cancellation adapter. |
sdk_v2/cpp/src/inferencing/session/oga_generator_cancellable.cc |
Implements OGA termination. |
sdk_v2/cpp/src/inferencing/session/live_session_registry.h |
Declares process-wide session tracking. |
sdk_v2/cpp/src/inferencing/session/live_session_registry.cc |
Implements session tracking. |
sdk_v2/cpp/src/inferencing/session/cancellable.h |
Defines the cancellable interface. |
sdk_v2/cpp/src/inferencing/generative/embeddings/embeddings_session.h |
Makes embedding generation cancellable. |
sdk_v2/cpp/src/inferencing/generative/embeddings/embeddings_session.cc |
Integrates deadlines into embedding inference. |
sdk_v2/cpp/src/inferencing/generative/chat/chat_session.cc |
Integrates cancellation into chat generation. |
sdk_v2/cpp/src/inferencing/generative/chat/chat_generator.h |
Makes chat generators cancellable. |
sdk_v2/cpp/src/inferencing/generative/audio/audio_session.cc |
Integrates cancellation into audio inference. |
sdk_v2/cpp/src/inferencing/generative/audio/audio_generator.h |
Makes audio generators cancellable. |
sdk_v2/cpp/src/c_api.cc |
Implements new C ABI functions. |
sdk_v2/cpp/include/foundry_local/foundry_local_cpp.inline.h |
Implements C++ wrapper methods. |
sdk_v2/cpp/include/foundry_local/foundry_local_cpp.h |
Declares public C++ APIs. |
sdk_v2/cpp/include/foundry_local/foundry_local_c.h |
Extends the public C ABI. |
sdk_v2/cpp/CMakeLists.txt |
Builds the new cancellation sources. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+124
to
+131
| std::vector<ICancellable*> generators; | ||
| std::vector<const Request*> requests; | ||
|
|
||
| { | ||
| std::lock_guard<std::mutex> lock(cancel_state_->mutex); | ||
| cancel_state_->cancel_requested = true; | ||
| generators = cancel_state_->active_generators; | ||
| requests = cancel_state_->active_requests_list; |
Comment on lines
+183
to
+185
| const bool woken = cancel_state_->cv.wait_for(lock, timeout, [this] { | ||
| return cancel_state_->active_requests == 0 || cancel_state_->cancel_requested; | ||
| }); |
Comment on lines
+69
to
+71
| void ArmDeadline() const { | ||
| canceled.store(false, std::memory_order_relaxed); | ||
| timed_out.store(false, std::memory_order_relaxed); |
Comment on lines
+24
to
+26
| std::vector<Session*> LiveSessionRegistry::Snapshot() const { | ||
| std::lock_guard<std::mutex> lock(mutex_); | ||
| return {sessions_.begin(), sessions_.end()}; |
Comment on lines
+12
to
+16
| try { | ||
| generator_.SetRuntimeOption("terminate_session", "1"); | ||
| } catch (const std::exception&) { | ||
| // SetRuntimeOption may not be supported by all ORT GenAI builds. | ||
| } |
Comment on lines
+286
to
+287
| // Detach before returning so the signal cannot cancel a subsequent reuse of | ||
| // this request, and so an long-lived signal does not retain it. |
Comment on lines
266
to
+270
| try { | ||
| ProcessRequestImpl(request, response); | ||
|
|
||
| tracker.SetStatus(ActionStatus::kSuccess); | ||
| tracker.SetStatus(request.canceled.load(std::memory_order_relaxed) ? ActionStatus::kCanceled | ||
| : ActionStatus::kSuccess); |
Comment on lines
+343
to
351
| if (!this.native.isDisposed()) { | ||
| try { | ||
| this.native.cancel(); | ||
| } catch { | ||
| // Best-effort: tear down regardless. | ||
| } | ||
| } | ||
|
|
||
| this.native.dispose(); |
Comment on lines
+307
to
+311
| // Also cancel at the session level. _activeStreamingCts covers only the streaming | ||
| // path; a non-streaming ProcessRequestAsync running on another thread is invisible | ||
| // to it, and it is exactly the case that otherwise pins the session and blocks | ||
| // model unload. | ||
| try { _session.Cancel(); } catch { } |
Comment on lines
+458
to
+464
| # Cancel at the session level rather than only cancelling the streaming request: | ||
| # a non-streaming generation on another thread is invisible to _stream_request, | ||
| # and it is exactly the case that otherwise pins the session and blocks unload. | ||
| try: | ||
| self.cancel() | ||
| except Exception: | ||
| pass |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Session::process_request()exposed no cancellation, timeout, or deadline mechanism.Request.cancel()was only wired into the streaming path, so a non-terminating non-streaming generation could permanently pin the session refcount.Consequences:
Model.Unload()failed with "1 session(s) still using it".FoundryLocalManager.close()exceeded its fixed 10-second drain deadline and left the model loaded.1073740791(0xC0000409).Fix
C++ core
ICancellableinterface andOgaGeneratorCancellableadapter soSessioncan interrupt ORT GenAI mid-compute (viaSetRuntimeOption("terminate_session", "1")), not just between token boundaries.Request::SetTimeout/ArmDeadline/ShouldStop— a re-armable wall-clock deadline that applies to streaming and non-streaming requests alike.Session::Cancel()with active-generator tracking and a deadline watchdog thread; wired into every chat/audio/embeddings generation loop.Cancel()now latchesrequest.canceledon every in-flight request (not just published generators), sofinish_reasonand history rollback are correct even when cancellation lands during prefill (before any generator is published).LiveSessionRegistrysoSessionManager::CancelAll()reaches every live session, including ones created via the direct API that never took aSessionRegistration— this is what fixes manager teardown and the 0xC0000409 crash.FOUNDRY_LOCAL_ERROR_TIMEOUTand two new C ABI vtable entries (Request_SetTimeoutMs,Session_Cancel), appended to preserve ABI ordering.Bindings (C++ wrapper, Python, C#, JS)
Request.SetTimeout()/set_timeout()/setTimeout(),Session.Cancel()/cancel().ProcessRequestAsync(C#) andprocessRequest(JS) now accept aCancellationToken/AbortSignalthat genuinely interrupts an in-flight generation (registered against the native cancel), not just prevents scheduling._close, C#Dispose, JSdispose) now cancels the session itself, covering the non-streaming case that was previously invisible to shutdown.Testing
Added
ChatSessionTestcoverage:FOUNDRY_LOCAL_ERROR_TIMEOUT.Request.Session::Cancel()stops an in-flight non-streaming request promptly.Cancel()before a request rejects immediately (no unbounded generation).Cancel()is idempotent and safe when idle.Full C++ unit suite: 1065/1069 passed (4 skips are pre-existing/unrelated audio integration tests requiring models not present in this environment).
Co-authored-by: Copilot App 223556219+Copilot@users.noreply.github.com