diff --git a/sdk_v2/cpp/CMakeLists.txt b/sdk_v2/cpp/CMakeLists.txt index d1020a1d5..9a6bf7b7e 100644 --- a/sdk_v2/cpp/CMakeLists.txt +++ b/sdk_v2/cpp/CMakeLists.txt @@ -215,6 +215,8 @@ set(FOUNDRY_LOCAL_SOURCES src/inferencing/generative/chat/chat_generator.cc src/inferencing/session/session.cc src/inferencing/session/session_manager.cc + src/inferencing/session/oga_generator_cancellable.cc + src/inferencing/session/live_session_registry.cc src/inferencing/generative/chat/chat_session.cc src/inferencing/generative/chat/chat_template.cc src/configuration.cc diff --git a/sdk_v2/cpp/include/foundry_local/foundry_local_c.h b/sdk_v2/cpp/include/foundry_local/foundry_local_c.h index fa1f63738..0e2e28bfd 100644 --- a/sdk_v2/cpp/include/foundry_local/foundry_local_c.h +++ b/sdk_v2/cpp/include/foundry_local/foundry_local_c.h @@ -184,6 +184,9 @@ typedef enum flErrorCode { FOUNDRY_LOCAL_ERROR_INVALID_USAGE = 4, FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED = 5, FOUNDRY_LOCAL_ERROR_NETWORK = 6, + /// An operation exceeded its caller-supplied deadline. Distinct from + /// FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED, which means an explicit cancel. + FOUNDRY_LOCAL_ERROR_TIMEOUT = 7, } flErrorCode; typedef enum flLogLevel { @@ -892,6 +895,24 @@ struct flInferenceApi { /// If all turns are undone, the cached generator is destroyed. FL_API_STATUS(Session_UndoTurns, _In_ flSession* session, size_t count); + /// Set a wall-clock budget for the request, in milliseconds. 0 (the default) means no + /// deadline. The clock starts when Session_ProcessRequest is called and is re-armed on + /// every call, so a Request may be reused. + /// + /// On expiry the in-flight generation is interrupted — including mid-compute, not just + /// at token boundaries — the session's model reference is released, and + /// Session_ProcessRequest returns FOUNDRY_LOCAL_ERROR_TIMEOUT. + FL_API_STATUS(Request_SetTimeoutMs, _In_ flRequest* request, uint64_t timeout_ms); + + /// Cancel the request currently in flight on this session, and make subsequent + /// Session_ProcessRequest calls fail fast with FOUNDRY_LOCAL_ERROR_INVALID_USAGE. + /// + /// Unlike Request_Cancel — which only reaches generation loops that poll the request — + /// this interrupts the underlying engine, so it also stops a non-streaming request that + /// is blocked inside a long prefill or decode. Safe to call from any thread; idempotent. + /// This is the supported way to release a session that is pinning a model. + FL_API_STATUS(Session_Cancel, _In_ flSession* session); + // End V1 }; diff --git a/sdk_v2/cpp/include/foundry_local/foundry_local_cpp.h b/sdk_v2/cpp/include/foundry_local/foundry_local_cpp.h index 30b69a529..747af8d34 100644 --- a/sdk_v2/cpp/include/foundry_local/foundry_local_cpp.h +++ b/sdk_v2/cpp/include/foundry_local/foundry_local_cpp.h @@ -20,6 +20,7 @@ #include "foundry_local/foundry_local_c.h" #include +#include #include #include #include @@ -955,6 +956,15 @@ class Request { /// Cancel the current request. Inferencing will stop as soon as possible. void Cancel(); + /// Set a wall-clock deadline for the request, in milliseconds. Pass 0 to disable. + /// + /// The deadline covers the whole ProcessRequest call, including prefill, and applies to + /// streaming and non-streaming alike. On expiry the run is interrupted and + /// ProcessRequest throws with FOUNDRY_LOCAL_ERROR_TIMEOUT. + /// + /// The deadline is re-armed on each ProcessRequest call, so a Request may be reused. + Request& SetTimeout(std::chrono::milliseconds timeout); + const flRequest* native_handle() const noexcept { return handle_.get(); } private: @@ -1001,6 +1011,17 @@ class Session { /// Populates the response with output items, finish reason, and usage. Response ProcessRequest(const Request& request); + /// Interrupt any requests currently running on this session. + /// + /// Thread-safe and callable from any thread — this is the point: it is meant to be + /// invoked from a different thread than the one blocked in ProcessRequest. It + /// interrupts inferencing mid-compute rather than only between tokens, so a + /// non-terminating generation cannot keep the session's reference to the model alive. + /// + /// The interrupted ProcessRequest returns promptly with a canceled finish reason. + /// Idempotent; safe to call when no request is running. + void Cancel(); + protected: detail::Base handle_; diff --git a/sdk_v2/cpp/include/foundry_local/foundry_local_cpp.inline.h b/sdk_v2/cpp/include/foundry_local/foundry_local_cpp.inline.h index 15c34b4d1..d3a57373b 100644 --- a/sdk_v2/cpp/include/foundry_local/foundry_local_cpp.inline.h +++ b/sdk_v2/cpp/include/foundry_local/foundry_local_cpp.inline.h @@ -1102,6 +1102,12 @@ inline void Request::Cancel() { Check(detail::inference_api()->Request_Cancel(handle_.get_mutable())); } +inline Request& Request::SetTimeout(std::chrono::milliseconds timeout) { + const auto ms = timeout.count() > 0 ? static_cast(timeout.count()) : uint64_t{0}; + Check(detail::inference_api()->Request_SetTimeoutMs(handle_.get_mutable(), ms)); + return *this; +} + inline flRequest* detail::CreateRequest() { flRequest* req = nullptr; Check(detail::inference_api()->Request_Create(&req)); @@ -1155,6 +1161,10 @@ inline Response Session::ProcessRequest(const Request& request) { return Response(response); } +inline void Session::Cancel() { + Check(detail::inference_api()->Session_Cancel(handle_.get_mutable())); +} + inline Session& Session::SetOptions(const RequestOptions& options) { KeyValuePairs kvp = detail::ToKeyValuePairs(options); Check(detail::inference_api()->Session_SetOptions(handle_.get_mutable(), kvp.native_handle())); diff --git a/sdk_v2/cpp/src/c_api.cc b/sdk_v2/cpp/src/c_api.cc index 6e89e534f..5e2ac8557 100644 --- a/sdk_v2/cpp/src/c_api.cc +++ b/sdk_v2/cpp/src/c_api.cc @@ -22,6 +22,7 @@ #include "manager.h" #include "ep_detection/ep_bootstrapper.h" +#include #include #include #include @@ -1625,6 +1626,17 @@ FL_API_STATUS_IMPL(Request_CancelImpl, flRequest* request) { API_IMPL_END } +FL_API_STATUS_IMPL(Request_SetTimeoutMsImpl, flRequest* request, uint64_t timeout_ms) { + API_IMPL_BEGIN + if (!request) { + return MakeStatus(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "null argument"); + } + + AsImpl(request)->SetTimeout(std::chrono::milliseconds(timeout_ms)); + return nullptr; + API_IMPL_END +} + // --- Response --- FL_API_STATUS_IMPL(Response_CreateImpl, flResponse** out_response) { @@ -1746,9 +1758,19 @@ FL_API_STATUS_IMPL(Session_SetOptionsImpl, flSession* session, const flKeyValueP API_IMPL_END } -FL_API_STATUS_IMPL(Session_ProcessRequestImpl, flSession* session, const flRequest* request, - flResponse** response) { +FL_API_STATUS_IMPL(Session_CancelImpl, flSession* session) { API_IMPL_BEGIN + if (!session) { + return MakeStatus(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "null session"); + } + + AsImpl(session)->Cancel(); + return nullptr; + API_IMPL_END +} + +FL_API_STATUS_IMPL(Session_ProcessRequestImpl, flSession* session, const flRequest* request, + flResponse** response) { API_IMPL_BEGIN if (!session || !request || !response) { return MakeStatus(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "null argument"); } @@ -1816,6 +1838,8 @@ static const flInferenceApi g_inference_api = { Session_RemoveToolDefinitionImpl, Session_GetTurnCountImpl, Session_UndoTurnsImpl, + Request_SetTimeoutMsImpl, + Session_CancelImpl, }; // ======================================================================== diff --git a/sdk_v2/cpp/src/inferencing/generative/audio/audio_generator.h b/sdk_v2/cpp/src/inferencing/generative/audio/audio_generator.h index 135938ef4..abbea4b58 100644 --- a/sdk_v2/cpp/src/inferencing/generative/audio/audio_generator.h +++ b/sdk_v2/cpp/src/inferencing/generative/audio/audio_generator.h @@ -2,6 +2,8 @@ // Licensed under the MIT License. #pragma once +#include "inferencing/session/cancellable.h" + #include namespace fl { @@ -10,7 +12,7 @@ namespace fl { /// One generator per request — not reusable, not thread-safe. /// Same pull-based iterator pattern as ChatGenerator: /// while (!IsDone()) { GenerateNextToken(); text += Decode(); } -class AudioGenerator { +class AudioGenerator : public ICancellable { public: virtual ~AudioGenerator() = default; @@ -39,7 +41,7 @@ class AudioGenerator { /// Request cancellation of generation. Thread-safe — can be called from another thread. /// After cancellation, IsDone() should return true on the next check. - virtual void Cancel() = 0; + void Cancel() override = 0; protected: AudioGenerator() = default; diff --git a/sdk_v2/cpp/src/inferencing/generative/audio/audio_session.cc b/sdk_v2/cpp/src/inferencing/generative/audio/audio_session.cc index e60ce8b14..a281f5434 100644 --- a/sdk_v2/cpp/src/inferencing/generative/audio/audio_session.cc +++ b/sdk_v2/cpp/src/inferencing/generative/audio/audio_session.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "inferencing/generative/audio/audio_session.h" +#include "inferencing/session/oga_generator_cancellable.h" #include "contracts/audio_transcriptions.h" #include "inferencing/generative/audio/onnx_audio_generator.h" @@ -248,22 +249,28 @@ void AudioSession::ProcessRequestImpl(const Request& request, Response& response std::vector> segments; segments.reserve(kInitialTokenCapacity); - while (!generator->IsDone() && !request.canceled) { - generator->GenerateNextToken(); - std::string token = generator->Decode(); + // Publish the generator so Session::Cancel() and the deadline watchdog can interrupt + // an in-flight compute, not just the loop between tokens. + { + ActiveGenerator active(*this, *generator); - if (!token.empty()) { - segments.push_back(MakeNoneSegment(token)); + while (!generator->IsDone() && !request.ShouldStop()) { + generator->GenerateNextToken(); + std::string token = generator->Decode(); - if (streaming_callback) { - streaming_callback->PushItem(MakeNoneSegment(token)); - } + if (!token.empty()) { + segments.push_back(MakeNoneSegment(token)); - token_texts.push_back(std::move(token)); - } + if (streaming_callback) { + streaming_callback->PushItem(MakeNoneSegment(token)); + } - if (request.canceled) { - generator->Cancel(); + token_texts.push_back(std::move(token)); + } + + if (request.canceled) { + generator->Cancel(); + } } } @@ -334,6 +341,11 @@ void AudioSession::ProcessStreamingAudio(const AudioItem& format_item, ItemQueue auto generator = OgaGenerator::Create(oga_model, *gen_params); auto tokenizer_stream = OgaTokenizerStream::Create(Model().Tokenizer().Oga()); + // Publish the raw generator so Session::Cancel() and the deadline watchdog can interrupt + // an in-flight encoder/decoder pass. Spans steps 3-5 below, which all drive this generator. + OgaGeneratorCancellable cancellable(*generator); + ActiveGenerator active(*this, cancellable); + auto streaming_callback = CreateCallbackHandler(request); std::vector token_texts; token_texts.reserve(kInitialTokenCapacity); @@ -352,7 +364,7 @@ void AudioSession::ProcessStreamingAudio(const AudioItem& format_item, ItemQueue } // 4. Read from queue until finished or cancelled - while (!request.canceled) { + while (!request.ShouldStop()) { auto item = queue.WaitAndPop(std::chrono::milliseconds(100)); if (!item) { @@ -428,7 +440,7 @@ void AudioSession::DecodeTokens(OgaGenerator& generator, OgaTokenizerStream& tok const std::unique_ptr& callback, const Request& request, int& completion_tokens) { - while (!generator.IsDone() && !generator.IsSessionTerminated() && !request.canceled) { + while (!generator.IsDone() && !generator.IsSessionTerminated() && !request.ShouldStop()) { generator.GenerateNextToken(); auto next_tokens = generator.GetNextTokens(); @@ -502,25 +514,31 @@ void AudioSession::ProcessAudioTranscriptionJson(const std::string& request_json // Generate token-by-token std::string text; - while (!generator->IsDone() && !original_request.canceled) { - generator->GenerateNextToken(); - std::string token = generator->Decode(); - - if (!token.empty()) { - text += token; - - if (is_streaming) { - // Emit streaming chunk as an OPENAI_JSON-tagged TextItem wrapping AudioTranscriptionResponse. - AudioTranscriptionResponse chunk; - chunk.id = response_id; - chunk.text = token; - streaming_callback->PushItem(std::make_unique(nlohmann::json(chunk).dump(), - FOUNDRY_LOCAL_TEXT_ITEM_TYPE_OPENAI_JSON)); + { + // Publish the generator so Session::Cancel() and the deadline watchdog can interrupt + // an in-flight compute, not just the loop between tokens. + ActiveGenerator active(*this, *generator); + + while (!generator->IsDone() && !original_request.ShouldStop()) { + generator->GenerateNextToken(); + std::string token = generator->Decode(); + + if (!token.empty()) { + text += token; + + if (is_streaming) { + // Emit streaming chunk as an OPENAI_JSON-tagged TextItem wrapping AudioTranscriptionResponse. + AudioTranscriptionResponse chunk; + chunk.id = response_id; + chunk.text = token; + streaming_callback->PushItem(std::make_unique(nlohmann::json(chunk).dump(), + FOUNDRY_LOCAL_TEXT_ITEM_TYPE_OPENAI_JSON)); + } } - } - if (original_request.canceled) { - generator->Cancel(); + if (original_request.canceled) { + generator->Cancel(); + } } } @@ -583,7 +601,7 @@ void AudioSession::DecodeNemotronTokens(OgaGenerator& generator, OgaTokenizerStr int& completion_tokens) const { const bool is_streaming = (streaming_callback != nullptr); - while (!generator.IsDone() && !generator.IsSessionTerminated() && !original_request.canceled) { + while (!generator.IsDone() && !generator.IsSessionTerminated() && !original_request.ShouldStop()) { generator.GenerateNextToken(); auto next_tokens = generator.GetNextTokens(); if (next_tokens.empty()) { @@ -618,7 +636,7 @@ void AudioSession::RunNemotronDecodePass(std::unique_ptr tensor const std::unique_ptr& streaming_callback, const std::string& response_id, const Request& original_request, int& completion_tokens) const { - if (!tensors || original_request.canceled) { + if (!tensors || original_request.ShouldStop()) { return; } @@ -659,6 +677,11 @@ void AudioSession::ProcessNemotronFileTranscription(const AudioTranscriptionRequ auto generator = OgaGenerator::Create(oga_model, *generator_params); TryNemotronLanguageId(*generator, language); + // Publish the raw generator so cancellation/deadline can interrupt an in-flight + // encoder or decode pass rather than only stopping between chunks. + OgaGeneratorCancellable cancellable(*generator); + ActiveGenerator active(*this, cancellable); + auto streaming_callback = CreateCallbackHandler(original_request); std::string response_id = ResponseConverter::GenerateId("audio"); @@ -667,7 +690,7 @@ void AudioSession::ProcessNemotronFileTranscription(const AudioTranscriptionRequ int completion_tokens = 0; constexpr size_t kNemotronSamplesPerChunk = 1600; // 100ms at 16kHz - for (size_t offset = 0; offset < samples.size() && !original_request.canceled; + for (size_t offset = 0; offset < samples.size() && !original_request.ShouldStop(); offset += kNemotronSamplesPerChunk) { size_t count = std::min(kNemotronSamplesPerChunk, samples.size() - offset); RunNemotronDecodePass(processor->Process(samples.data() + offset, count), *generator, *tokenizer_stream, text, diff --git a/sdk_v2/cpp/src/inferencing/generative/chat/chat_generator.h b/sdk_v2/cpp/src/inferencing/generative/chat/chat_generator.h index 77430d984..3f3f1237e 100644 --- a/sdk_v2/cpp/src/inferencing/generative/chat/chat_generator.h +++ b/sdk_v2/cpp/src/inferencing/generative/chat/chat_generator.h @@ -2,6 +2,8 @@ // Licensed under the MIT License. #pragma once +#include "inferencing/session/cancellable.h" + #include namespace fl { @@ -10,13 +12,12 @@ namespace fl { /// One generator per request — not reusable, not thread-safe. /// Follows the classic pull-based iterator pattern: /// while (!IsDone()) { GenerateNextToken(); text += Decode(); } -class ChatGenerator { +class ChatGenerator : public ICancellable { public: virtual ~ChatGenerator() = default; ChatGenerator(const ChatGenerator&) = delete; ChatGenerator& operator=(const ChatGenerator&) = delete; - /// Returns true when generation is complete (EOS token, max_length, or stop condition). virtual bool IsDone() const = 0; @@ -39,7 +40,7 @@ class ChatGenerator { /// Request cancellation of generation. Thread-safe — can be called from another thread. /// After cancellation, IsDone() should return true on the next check. - virtual void Cancel() = 0; + void Cancel() override = 0; protected: ChatGenerator() = default; diff --git a/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.cc b/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.cc index eb76ef94d..2db59de3d 100644 --- a/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.cc +++ b/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.cc @@ -574,20 +574,27 @@ void ChatSession::ProcessRequestImpl(const Request& request, Response& response) } }; - while (!cached_generator_->IsDone() && !request.canceled) { - cached_generator_->GenerateNextToken(); - std::string token = cached_generator_->Decode(); - ++output_tokens; + // Publish the generator so Session::Cancel() and the deadline watchdog can interrupt + // it mid-compute, not just between tokens. Scoped to the generation loop: the cached + // generator may be reset below, and it must not stay published past this point. + { + ActiveGenerator active(*this, *cached_generator_); - if (!token.empty()) { - text += token; - emit_segments(splitter.Push(token)); - } + while (!cached_generator_->IsDone() && !request.ShouldStop()) { + cached_generator_->GenerateNextToken(); + std::string token = cached_generator_->Decode(); + ++output_tokens; - // Enforce max_output_tokens — with use_full_context the OGA max_length - // is the entire context window, so we must cap output ourselves. - if (max_output > 0 && output_tokens >= max_output) { - break; + if (!token.empty()) { + text += token; + emit_segments(splitter.Push(token)); + } + + // Enforce max_output_tokens — with use_full_context the OGA max_length + // is the entire context window, so we must cap output ourselves. + if (max_output > 0 && output_tokens >= max_output) { + break; + } } } @@ -794,13 +801,19 @@ void ChatSession::ProcessChatCompletionsJson(const std::string& request_json, co // Generate token-by-token std::string text; - while (!generator->IsDone() && !original_request.canceled) { - generator->GenerateNextToken(); - std::string token = generator->Decode(); + { + // Publish the generator so Session::Cancel() and the deadline watchdog can interrupt + // an in-flight compute, not just the loop between tokens. + ActiveGenerator active(*this, *generator); - if (!token.empty()) { - text += token; - process_segments(splitter.Push(token)); + while (!generator->IsDone() && !original_request.ShouldStop()) { + generator->GenerateNextToken(); + std::string token = generator->Decode(); + + if (!token.empty()) { + text += token; + process_segments(splitter.Push(token)); + } } } diff --git a/sdk_v2/cpp/src/inferencing/generative/embeddings/embeddings_session.cc b/sdk_v2/cpp/src/inferencing/generative/embeddings/embeddings_session.cc index dfa1c512b..8d9e453ba 100644 --- a/sdk_v2/cpp/src/inferencing/generative/embeddings/embeddings_session.cc +++ b/sdk_v2/cpp/src/inferencing/generative/embeddings/embeddings_session.cc @@ -4,6 +4,7 @@ #include "contracts/embeddings.h" #include "exception.h" +#include "inferencing/session/oga_generator_cancellable.h" #include "inferencing/generative/embeddings/fp16.h" #include "inferencing/generative/genai_model_instance.h" #include "items/tensor_item.h" @@ -73,7 +74,14 @@ void EmbeddingsSession::ProcessRequestImpl(const Request& request, Response& res } // Single batched forward pass for all inputs. - auto embeddings = GenerateEmbeddingsBatch(inputs); + auto embeddings = GenerateEmbeddingsBatch(inputs, request); + + // A cancelled or timed-out batch is incomplete — emit no partial items and let the + // caller see it stopped early, rather than silently returning fewer vectors than inputs. + if (request.canceled) { + response.finish_reason = FOUNDRY_LOCAL_FINISH_NONE; + return; + } // Wrap each embedding as a TensorItem in the response. The vector is heap-allocated // and ownership is transferred to the TensorItem deleter via deleter_user_data_, so @@ -95,7 +103,7 @@ void EmbeddingsSession::ProcessRequestImpl(const Request& request, Response& res } void EmbeddingsSession::ProcessEmbeddingsJson(const std::string& request_json, - const Request& /*original_request*/, + const Request& original_request, Response& response) { // Parse the OpenAI embeddings request. Let nlohmann::json::parse_error propagate — // matches AudioSession::ProcessAudioTranscriptionJson behavior. @@ -117,7 +125,14 @@ void EmbeddingsSession::ProcessEmbeddingsJson(const std::string& request_json, // Reuse GenerateEmbeddingsBatch so the typed and JSON paths stay bit-for-bit equal // under future refactors (parity test relies on this). if (!inputs.empty()) { - auto embeddings = GenerateEmbeddingsBatch(inputs); + auto embeddings = GenerateEmbeddingsBatch(inputs, original_request); + + // Cancelled or timed out mid-batch: report the stop instead of returning a + // short data array that a client would read as a complete result. + if (original_request.canceled) { + response.finish_reason = FOUNDRY_LOCAL_FINISH_NONE; + return; + } output.data.reserve(embeddings.size()); for (size_t i = 0; i < embeddings.size(); ++i) { @@ -142,7 +157,7 @@ void EmbeddingsSession::ProcessEmbeddingsJson(const std::string& request_json, } std::vector> EmbeddingsSession::GenerateEmbeddingsBatch( - const std::vector& inputs) { + const std::vector& inputs, const Request& request) { // Process each input independently (batch_size=1 per forward pass). // // Embedding models like Qwen3-Embedding use bidirectional attention — @@ -157,16 +172,22 @@ std::vector> EmbeddingsSession::GenerateEmbeddingsBatch( results.reserve(inputs.size()); for (const auto& input : inputs) { - results.push_back(GenerateSingleEmbedding(input)); + // Check between inputs: a large batch is otherwise uninterruptible, which is what + // lets a runaway embeddings request pin the session refcount through shutdown. + if (request.ShouldStop()) { + break; + } + + results.push_back(GenerateSingleEmbedding(input, request)); } logger_.Log(LogLevel::Verbose, - fmt::format("Embeddings: processed {} input(s)", inputs.size())); + fmt::format("Embeddings: processed {} of {} input(s)", results.size(), inputs.size())); return results; } -std::vector EmbeddingsSession::GenerateSingleEmbedding(const std::string& input) { +std::vector EmbeddingsSession::GenerateSingleEmbedding(const std::string& input, const Request& request) { auto& oga_model = model_.GetOgaModel(); // 1. Tokenize and append EOS. Encode is serialized on the model's shared tokenizer. @@ -186,9 +207,20 @@ std::vector EmbeddingsSession::GenerateSingleEmbedding(const std::string& auto generator = OgaGenerator::Create(oga_model, *gen_params); generator->AppendTokenSequences(*sequences); + // Publish the generator so cancellation or an expired deadline interrupts this forward + // pass. A single long input can exceed the budget without ever reaching the loop check + // in GenerateEmbeddingsBatch. + OgaGeneratorCancellable cancellable(*generator); + ActiveGenerator active(*this, cancellable); + // 3. Single forward pass. generator->GenerateNextToken(); + // The pass may have been terminated mid-compute, leaving hidden_states unusable. + if (request.ShouldStop()) { + return {}; + } + // 4. Extract hidden_states. Shape: [1, token_count, hidden_size] auto hidden_states = generator->GetOutput("hidden_states"); auto shape = hidden_states->Shape(); diff --git a/sdk_v2/cpp/src/inferencing/generative/embeddings/embeddings_session.h b/sdk_v2/cpp/src/inferencing/generative/embeddings/embeddings_session.h index 26f9a317d..e5fd60d72 100644 --- a/sdk_v2/cpp/src/inferencing/generative/embeddings/embeddings_session.h +++ b/sdk_v2/cpp/src/inferencing/generative/embeddings/embeddings_session.h @@ -37,10 +37,16 @@ class EmbeddingsSession : public Session { /// Generate L2-normalized embedding vectors for a list of inputs. /// Each input is processed independently (batch_size=1) to avoid /// padding artifacts with bidirectional-attention embedding models. - std::vector> GenerateEmbeddingsBatch(const std::vector& inputs); + /// + /// Stops early if `request` is cancelled or its deadline expires, so the returned + /// vector may be shorter than `inputs`. Callers must check the request state before + /// treating the result as complete. + std::vector> GenerateEmbeddingsBatch(const std::vector& inputs, + const Request& request); /// Generate a single L2-normalized embedding vector for one input string. - std::vector GenerateSingleEmbedding(const std::string& input); + /// Returns an empty vector if the forward pass was interrupted by cancellation or a timeout. + std::vector GenerateSingleEmbedding(const std::string& input, const Request& request); /// Process a request whose first item is a TEXT item tagged OPENAI_JSON containing an /// OpenAI EmbeddingCreateRequest payload. Parses the JSON, runs generation via the diff --git a/sdk_v2/cpp/src/inferencing/session/cancellable.h b/sdk_v2/cpp/src/inferencing/session/cancellable.h new file mode 100644 index 000000000..48fe194f0 --- /dev/null +++ b/sdk_v2/cpp/src/inferencing/session/cancellable.h @@ -0,0 +1,24 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#pragma once + +namespace fl { + +/// Anything that can be interrupted from a thread other than the one driving it. +/// +/// Implemented by the per-request generators (chat, audio). A generation loop only +/// observes `Request::ShouldStop()` between tokens, which is not enough on its own: +/// a long prefill or a single slow decode step can block inside the ORT GenAI engine +/// for an unbounded time. Cancel() reaches into the engine (terminate_session) so an +/// external canceller or an expired deadline interrupts mid-compute. +/// +/// Cancel() must be safe to call from any thread, at any point in the generator's +/// lifetime, and must be idempotent. +class ICancellable { + public: + virtual ~ICancellable() = default; + + virtual void Cancel() = 0; +}; + +} // namespace fl diff --git a/sdk_v2/cpp/src/inferencing/session/live_session_registry.cc b/sdk_v2/cpp/src/inferencing/session/live_session_registry.cc new file mode 100644 index 000000000..3aff08a85 --- /dev/null +++ b/sdk_v2/cpp/src/inferencing/session/live_session_registry.cc @@ -0,0 +1,29 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#include "inferencing/session/live_session_registry.h" + +namespace fl { + +LiveSessionRegistry& LiveSessionRegistry::Instance() { + // Leaked intentionally: sessions may be destroyed during static destruction, and a + // destroyed registry would make Remove() a use-after-free. + static LiveSessionRegistry* instance = new LiveSessionRegistry(); + return *instance; +} + +void LiveSessionRegistry::Add(Session& session) { + std::lock_guard lock(mutex_); + sessions_.insert(&session); +} + +void LiveSessionRegistry::Remove(Session& session) { + std::lock_guard lock(mutex_); + sessions_.erase(&session); +} + +std::vector LiveSessionRegistry::Snapshot() const { + std::lock_guard lock(mutex_); + return {sessions_.begin(), sessions_.end()}; +} + +} // namespace fl diff --git a/sdk_v2/cpp/src/inferencing/session/live_session_registry.h b/sdk_v2/cpp/src/inferencing/session/live_session_registry.h new file mode 100644 index 000000000..513297192 --- /dev/null +++ b/sdk_v2/cpp/src/inferencing/session/live_session_registry.h @@ -0,0 +1,43 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#pragma once + +#include +#include +#include + +namespace fl { + +class Session; + +/// Process-wide set of every live Session, maintained by Session's constructor and +/// destructor. +/// +/// SessionManager only knows about sessions that took a SessionRegistration, which the +/// web-service handlers do but the direct API does not. Cancellation must reach *every* +/// session: a non-terminating direct-API request is precisely the case that pins the +/// model refcount and stalls manager teardown. +/// +/// This tracks sessions for cancellation only. It deliberately has no bearing on +/// SessionManager::WaitForDrain(), because an idle session the caller still owns must +/// not block shutdown. +class LiveSessionRegistry { + public: + static LiveSessionRegistry& Instance(); + + void Add(Session& session); + void Remove(Session& session); + + /// Snapshot of the live sessions. Returned by value so callers can cancel without + /// holding the lock — Session::Cancel() reaches into the inference engine and would + /// otherwise deadlock against a session unwinding and calling Remove(). + std::vector Snapshot() const; + + private: + LiveSessionRegistry() = default; + + mutable std::mutex mutex_; + std::unordered_set sessions_; +}; + +} // namespace fl diff --git a/sdk_v2/cpp/src/inferencing/session/oga_generator_cancellable.cc b/sdk_v2/cpp/src/inferencing/session/oga_generator_cancellable.cc new file mode 100644 index 000000000..1974d6d35 --- /dev/null +++ b/sdk_v2/cpp/src/inferencing/session/oga_generator_cancellable.cc @@ -0,0 +1,19 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#include "inferencing/session/oga_generator_cancellable.h" + +#include + +namespace fl { + +void OgaGeneratorCancellable::Cancel() { + // Engine-level termination interrupts an in-flight compute (e.g. a long encoder pass), + // which a between-token flag check cannot do. Mirrors OnnxAudioGenerator::Cancel. + try { + generator_.SetRuntimeOption("terminate_session", "1"); + } catch (const std::exception&) { + // SetRuntimeOption may not be supported by all ORT GenAI builds. + } +} + +} // namespace fl diff --git a/sdk_v2/cpp/src/inferencing/session/oga_generator_cancellable.h b/sdk_v2/cpp/src/inferencing/session/oga_generator_cancellable.h new file mode 100644 index 000000000..dc1262f5c --- /dev/null +++ b/sdk_v2/cpp/src/inferencing/session/oga_generator_cancellable.h @@ -0,0 +1,31 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#pragma once + +#include "inferencing/session/cancellable.h" + +#include + +struct OgaGenerator; + +namespace fl { + +/// Adapts a raw OgaGenerator to ICancellable. +/// +/// Some paths (streaming PCM transcription, Nemotron decode) drive an OgaGenerator +/// directly instead of going through OnnxAudioGenerator, so they have no Cancel() of +/// their own. Wrapping the generator lets Session publish it for cancellation, which +/// is what makes those loops interruptible mid-compute rather than only between tokens. +/// +/// Non-owning: the wrapped generator must outlive this adapter. +class OgaGeneratorCancellable : public ICancellable { + public: + explicit OgaGeneratorCancellable(OgaGenerator& generator) : generator_(generator) {} + + void Cancel() override; + + private: + OgaGenerator& generator_; +}; + +} // namespace fl diff --git a/sdk_v2/cpp/src/inferencing/session/request.h b/sdk_v2/cpp/src/inferencing/session/request.h index f13957eef..fba33ca1b 100644 --- a/sdk_v2/cpp/src/inferencing/session/request.h +++ b/sdk_v2/cpp/src/inferencing/session/request.h @@ -6,6 +6,8 @@ #include "util/key_value_pairs.h" #include +#include +#include #include #include @@ -23,18 +25,28 @@ struct Request { /// Uses relaxed ordering since it is a one-way flag and exact timing doesn't matter. mutable std::atomic canceled{false}; + /// Set to true when the request was stopped because its deadline expired rather than + /// by an explicit Cancel(). Callers use this to distinguish a timeout from a user cancel. + mutable std::atomic timed_out{false}; + Request() = default; Request(Request&& other) noexcept : items(std::move(other.items)), options(std::move(other.options)), canceled(other.canceled.load(std::memory_order_relaxed)), + timed_out(other.timed_out.load(std::memory_order_relaxed)), + timeout_ms_(other.timeout_ms_.load(std::memory_order_relaxed)), + deadline_ticks_(other.deadline_ticks_.load(std::memory_order_relaxed)), owned_items(std::move(other.owned_items)) {} Request& operator=(Request&& other) noexcept { items = std::move(other.items); options = std::move(other.options); canceled.store(other.canceled.load(std::memory_order_relaxed), std::memory_order_relaxed); + timed_out.store(other.timed_out.load(std::memory_order_relaxed), std::memory_order_relaxed); + timeout_ms_.store(other.timeout_ms_.load(std::memory_order_relaxed), std::memory_order_relaxed); + deadline_ticks_.store(other.deadline_ticks_.load(std::memory_order_relaxed), std::memory_order_relaxed); owned_items = std::move(other.owned_items); return *this; } @@ -42,6 +54,57 @@ struct Request { Request(const Request&) = delete; Request& operator=(const Request&) = delete; + /// Set a wall-clock budget for the request. Zero (the default) means no deadline. + /// Takes effect on the next ArmDeadline() — i.e. the next ProcessRequest call. + void SetTimeout(std::chrono::milliseconds timeout) { + timeout_ms_.store(timeout.count() < 0 ? 0 : static_cast(timeout.count()), std::memory_order_relaxed); + } + + std::chrono::milliseconds Timeout() const { + return std::chrono::milliseconds(timeout_ms_.load(std::memory_order_relaxed)); + } + + /// Start the timeout clock and clear per-run stop state. Called by Session::ProcessRequest + /// so a Request reused across calls gets a fresh budget each time. + void ArmDeadline() const { + canceled.store(false, std::memory_order_relaxed); + timed_out.store(false, std::memory_order_relaxed); + + const auto budget = timeout_ms_.load(std::memory_order_relaxed); + if (budget == 0) { + deadline_ticks_.store(0, std::memory_order_relaxed); + return; + } + + const auto deadline = Clock::now() + std::chrono::milliseconds(budget); + deadline_ticks_.store(static_cast(deadline.time_since_epoch().count()), std::memory_order_relaxed); + } + + /// Clear the deadline so a later check cannot trip after the run has ended. + void DisarmDeadline() const { deadline_ticks_.store(0, std::memory_order_relaxed); } + + /// True once the request should stop — either explicitly cancelled or past its deadline. + /// Deadline expiry latches `canceled` so the existing post-loop cancellation handling + /// (rewind, finish_reason, history rollback) applies unchanged to timeouts. + bool ShouldStop() const { + if (canceled.load(std::memory_order_relaxed)) { + return true; + } + + const auto ticks = deadline_ticks_.load(std::memory_order_relaxed); + if (ticks == 0) { + return false; + } + + if (Clock::now().time_since_epoch().count() < ticks) { + return false; + } + + timed_out.store(true, std::memory_order_relaxed); + canceled.store(true, std::memory_order_relaxed); + return true; + } + /// Add a pre-allocated owned item. void AddOwnedItem(std::unique_ptr item) { items.push_back(item.get()); @@ -54,6 +117,11 @@ struct Request { } private: + using Clock = std::chrono::steady_clock; + + mutable std::atomic timeout_ms_{0}; + /// steady_clock time_point ticks for the current run's deadline; 0 means unarmed. + mutable std::atomic deadline_ticks_{0}; std::vector> owned_items; // owned items (lifetime) }; diff --git a/sdk_v2/cpp/src/inferencing/session/session.cc b/sdk_v2/cpp/src/inferencing/session/session.cc index e8f90ee69..8e2fa06ab 100644 --- a/sdk_v2/cpp/src/inferencing/session/session.cc +++ b/sdk_v2/cpp/src/inferencing/session/session.cc @@ -7,6 +7,7 @@ #include "inferencing/generative/chat/chat_session.h" #include "inferencing/generative/embeddings/embeddings_session.h" #include "inferencing/model_load_manager.h" +#include "inferencing/session/live_session_registry.h" #include "inferencing/session/session_manager.h" #include "manager.h" #include "model.h" @@ -26,9 +27,34 @@ Session::Session(const fl::Model& catalog_model, ILogger& logger, ITelemetry& te logger_(logger), telemetry_(telemetry), allow_concurrent_requests_(allow_concurrent_requests) { + LiveSessionRegistry::Instance().Add(*this); } -Session::~Session() = default; +// Moving relocates the session, so the registry must follow the object's address. +// The moved-from shell stays registered until its own destructor runs; cancelling it is +// harmless because its cancel state moved with it. +Session::Session(Session&& other) noexcept + : catalog_model_(other.catalog_model_), + logger_(other.logger_), + telemetry_(other.telemetry_), + tool_definitions_(std::move(other.tool_definitions_)), + session_options_(std::move(other.session_options_)), + callback_fn_(std::move(other.callback_fn_)), + callback_user_data_(other.callback_user_data_), + allow_concurrent_requests_(other.allow_concurrent_requests_), + request_mutex_(std::move(other.request_mutex_)), + cancel_state_(std::move(other.cancel_state_)) { + // Give the moved-from shell fresh state rather than null pointers: it stays live (and + // reachable from the registry) until its destructor runs, and Cancel() may race with it. + other.request_mutex_ = std::make_unique(); + other.cancel_state_ = std::make_unique(); + + LiveSessionRegistry::Instance().Add(*this); +} + +Session::~Session() { + LiveSessionRegistry::Instance().Remove(*this); +} std::unique_ptr Session::Create(const fl::Model& model) { auto& mgr = Manager::Instance(); @@ -94,6 +120,91 @@ void Session::AddToolDefinition(ToolDefinition tool_def) { tool_definitions_.push_back(std::move(tool_def)); } +void Session::Cancel() { + std::vector generators; + std::vector requests; + + { + std::lock_guard lock(cancel_state_->mutex); + cancel_state_->cancel_requested = true; + generators = cancel_state_->active_generators; + requests = cancel_state_->active_requests_list; + } + + // Latch the flag on every in-flight request, not just the generators. Interrupting the + // engine alone is not enough: the loops would exit but the run would still be reported + // as a natural stop, and cancelled turns would be committed to history. Cancellation + // during prefill is exactly this case — it produces no tokens, so nothing else marks it. + for (const auto* request : requests) { + request->canceled.store(true, std::memory_order_relaxed); + } + + // Wake the deadline watchdog so it exits instead of sleeping out the full budget. + cancel_state_->cv.notify_all(); + + // Cancel outside the lock: each Cancel() reaches into ORT GenAI, and holding the + // mutex across that would block the request thread's generator bookkeeping. + for (auto* generator : generators) { + generator->Cancel(); + } +} + +void Session::AddActiveGenerator(ICancellable* generator) { + bool cancel_now = false; + + { + std::lock_guard lock(cancel_state_->mutex); + cancel_state_->active_generators.push_back(generator); + + // Cancel() may have landed between the caller's pre-flight check and this + // publication. Apply the pending stop to the newly-visible generator so the + // request cannot slip past an already-issued cancellation. + cancel_now = cancel_state_->cancel_requested; + } + + if (cancel_now) { + generator->Cancel(); + } +} + +void Session::RemoveActiveGenerator(ICancellable* generator) { + std::lock_guard lock(cancel_state_->mutex); + auto& generators = cancel_state_->active_generators; + generators.erase(std::remove(generators.begin(), generators.end(), generator), generators.end()); +} + +void Session::WatchDeadline(const Request& request) { + const auto timeout = request.Timeout(); + + std::unique_lock lock(cancel_state_->mutex); + + // Wait out the budget, but wake early if the request finished or was cancelled — + // otherwise ProcessRequest would block on joining this thread for the full timeout. + const bool woken = cancel_state_->cv.wait_for(lock, timeout, [this] { + return cancel_state_->active_requests == 0 || cancel_state_->cancel_requested; + }); + + if (woken) { + return; + } + + // Deadline expired. Latch the timeout on the request so generation loops stop at the + // next token boundary, and cancel the active generators to interrupt an in-flight + // compute (a long prefill can exceed the budget without ever reaching a boundary). + request.canceled.store(true, std::memory_order_relaxed); + request.timed_out.store(true, std::memory_order_relaxed); + + auto generators = cancel_state_->active_generators; + lock.unlock(); + + logger_.Log(LogLevel::Warning, + fmt::format("request exceeded its {}ms deadline; cancelling", timeout.count())); + + for (auto* generator : generators) { + generator->Cancel(); + } +} + void Session::ProcessRequest(const Request& request, Response& response) { // Serialize requests unless the derived class opted into concurrency. std::unique_lock lock(*request_mutex_, std::defer_lock); @@ -101,17 +212,72 @@ void Session::ProcessRequest(const Request& request, Response& response) { lock.lock(); } + // A session cancelled during teardown must not start new work — otherwise a caller + // looping over requests could keep the model refcount pinned past Manager::Shutdown. + { + std::lock_guard cancel_lock(cancel_state_->mutex); + if (cancel_state_->cancel_requested) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_USAGE, "session has been cancelled"); + } + + ++cancel_state_->active_requests; + cancel_state_->active_requests_list.push_back(&request); + } + + // Start the wall-clock budget (if any) and clear stale stop state from a prior run. + request.ArmDeadline(); + + // Watchdog enforces the deadline for paths that would otherwise block indefinitely + // inside the engine. Only started when a timeout was requested — no cost otherwise. + std::thread watchdog; + if (request.Timeout().count() > 0) { + watchdog = std::thread(&Session::WatchDeadline, this, std::cref(request)); + } + + // Guarantees the watchdog is woken and joined on every exit path, including throws. + // Declared after the thread so it runs first on unwind. + struct RunScope { + Session& session; + const Request& request; + std::thread& watchdog; + + ~RunScope() { + { + std::lock_guard lock(session.cancel_state_->mutex); + --session.cancel_state_->active_requests; + + auto& list = session.cancel_state_->active_requests_list; + list.erase(std::remove(list.begin(), list.end(), &request), list.end()); + } + + session.cancel_state_->cv.notify_all(); + + if (watchdog.joinable()) { + watchdog.join(); + } + + request.DisarmDeadline(); + } + } run_scope{*this, request, watchdog}; + ActionTracker tracker(Action::kSessionProcessRequest, telemetry_); tracker.SetModelId(CatalogModel().Id()); try { ProcessRequestImpl(request, response); - tracker.SetStatus(ActionStatus::kSuccess); + tracker.SetStatus(request.canceled.load(std::memory_order_relaxed) ? ActionStatus::kCanceled + : ActionStatus::kSuccess); } catch (const std::exception& ex) { tracker.RecordException(ex); throw; } + + // A timeout is a failure of the caller's contract, not a silent truncation: surface it + // so callers can distinguish "the model stopped early" from "we ran out of time". + if (request.timed_out.load(std::memory_order_relaxed)) { + FL_THROW(FOUNDRY_LOCAL_ERROR_TIMEOUT, "request timed out after ", request.Timeout().count(), "ms"); + } } } // namespace fl diff --git a/sdk_v2/cpp/src/inferencing/session/session.h b/sdk_v2/cpp/src/inferencing/session/session.h index f3ec0e5f1..9cf4b14b9 100644 --- a/sdk_v2/cpp/src/inferencing/session/session.h +++ b/sdk_v2/cpp/src/inferencing/session/session.h @@ -3,15 +3,19 @@ #pragma once #include +#include +#include #include #include #include #include +#include #include #include #include "inferencing/session/callback_handler.h" +#include "inferencing/session/cancellable.h" #include "inferencing/session/request.h" #include "inferencing/session/response.h" #include "inferencing/session/types.h" @@ -34,7 +38,7 @@ class Session { public: virtual ~Session(); - Session(Session&&) = default; + Session(Session&&) noexcept; Session& operator=(Session&&) = delete; Session(const Session&) = delete; @@ -50,8 +54,21 @@ class Session { /// class, then waits for all async streaming callbacks to complete. /// Waiting here keeps the Request reference valid for the lifetime of any /// in-flight callbacks and ensures the Response is fully populated on return. + /// + /// Cancellation applies to every path, streaming or not: + /// - `Session::Cancel()` from another thread stops the run promptly. + /// - `Request::SetTimeout()` bounds the run with a wall-clock deadline. + /// Both work by flagging the Request and calling Cancel() on the active generator, + /// so an in-flight ORT GenAI compute is interrupted rather than waited out. void ProcessRequest(const Request& request, Response& response); + /// Signal the in-flight request (if any) to stop, and cause the next ProcessRequest + /// on this session to abort immediately. Safe to call from any thread; idempotent. + /// + /// This is the teardown hook: it guarantees a runaway non-streaming generation + /// releases the session's model refcount instead of pinning the model loaded. + void Cancel(); + /// Add a tool definition to this session. /// @throws fl::Exception if tool_def.json_schema is not valid JSON. void AddToolDefinition(ToolDefinition tool_def); @@ -143,7 +160,46 @@ class Session { const KeyValuePairs& SessionOptions() const { return session_options_; } + /// RAII guard publishing the generator currently driving a request so that + /// Session::Cancel() and the deadline watchdog can interrupt it mid-compute. + /// + /// Derived classes create one for the scope in which a generator is live. If + /// cancellation was already requested when the guard is constructed, the + /// generator is cancelled immediately — this closes the race where Cancel() + /// lands between the pre-flight check and the generator becoming visible. + /// + /// Multiple guards may be live at once: sessions that opt into concurrency + /// (embeddings) run several requests in parallel, and a nested scope can publish + /// a second generator. All registered generators are cancelled together. + class ActiveGenerator { + public: + ActiveGenerator(Session& session, ICancellable& generator) + : session_(session), generator_(generator) { + session_.AddActiveGenerator(&generator_); + } + + ~ActiveGenerator() { session_.RemoveActiveGenerator(&generator_); } + + ActiveGenerator(const ActiveGenerator&) = delete; + ActiveGenerator& operator=(const ActiveGenerator&) = delete; + + private: + Session& session_; + ICancellable& generator_; + }; + private: + /// Publish a generator driving a request. Cancels it inline if a stop was already + /// requested, so a generator created after Cancel() cannot run unbounded. + void AddActiveGenerator(ICancellable* generator); + + /// Withdraw a generator once its scope ends. + void RemoveActiveGenerator(ICancellable* generator); + + /// Body of the deadline watchdog thread. Sleeps until the request's deadline and + /// then interrupts the run, unless woken earlier because the request completed. + void WatchDeadline(const Request& request); + const fl::Model& catalog_model_; ILogger& logger_; ITelemetry& telemetry_; @@ -153,6 +209,26 @@ class Session { void* callback_user_data_ = nullptr; const bool allow_concurrent_requests_; mutable std::unique_ptr request_mutex_ = std::make_unique(); + + /// Guards active_generator_/cancel_requested_ and pairs with the condition variable for + /// the watchdog. Held behind a unique_ptr because Session must remain movable (the + /// Responses API caches ChatSessions by move) and mutex/condition_variable are not. + /// Separate from request_mutex_: Cancel() must be serviceable while a request holds that lock. + struct CancelState { + std::mutex mutex; + std::condition_variable cv; + std::vector active_generators; + /// In-flight requests, so Cancel() can latch the flag that drives finish_reason and + /// history rollback. Tracked alongside generators because a request may be cancelled + /// before it publishes one (e.g. during prefill). + std::vector active_requests_list; + bool cancel_requested = false; + /// Number of requests currently running, so the watchdog knows when to stop waiting. + /// A count rather than a flag because concurrent sessions overlap requests. + int active_requests = 0; + }; + + std::unique_ptr cancel_state_ = std::make_unique(); }; } // namespace fl diff --git a/sdk_v2/cpp/src/inferencing/session/session_manager.cc b/sdk_v2/cpp/src/inferencing/session/session_manager.cc index a48e81bb6..a480e5981 100644 --- a/sdk_v2/cpp/src/inferencing/session/session_manager.cc +++ b/sdk_v2/cpp/src/inferencing/session/session_manager.cc @@ -4,8 +4,12 @@ #include "exception.h" #include "inferencing/generative/chat/chat_session.h" +#include "inferencing/session/live_session_registry.h" +#include "inferencing/session/session.h" #include +#include + #include namespace fl { @@ -53,11 +57,25 @@ void SessionManager::CancelAll() { // Clear cache — frees idle cached sessions so they don't block drain. ClearCache(); - std::lock_guard lock(mutex_); - logger_.Log(LogLevel::Information, - fmt::format("SessionManager: cancelling all sessions ({} active)", sessions_.size())); + // Cancel every live session, not just registered ones: direct-API sessions never take a + // SessionRegistration, yet a runaway request on one is exactly what pins the model. + // + // Snapshot first and cancel outside any lock — Session::Cancel() reaches into the ORT + // GenAI engine, and a cancelled session unwinding calls back into Deregister(). + std::vector to_cancel = LiveSessionRegistry::Instance().Snapshot(); - // Future (Phase 3): iterate sessions_ and call Cancel() on each + logger_.Log(LogLevel::Information, + fmt::format("SessionManager: cancelling all sessions ({} live)", to_cancel.size())); + + for (auto* session : to_cancel) { + try { + session->Cancel(); + } catch (const std::exception& ex) { + // Cancellation is best-effort during shutdown; one wedged session must not + // prevent the others from being signalled. + logger_.Log(LogLevel::Warning, fmt::format("SessionManager: failed to cancel a session: {}", ex.what())); + } + } } void SessionManager::WaitForDrain(std::chrono::milliseconds timeout) { diff --git a/sdk_v2/cpp/src/inferencing/session/session_manager.h b/sdk_v2/cpp/src/inferencing/session/session_manager.h index 121a24323..08edcb16d 100644 --- a/sdk_v2/cpp/src/inferencing/session/session_manager.h +++ b/sdk_v2/cpp/src/inferencing/session/session_manager.h @@ -62,6 +62,10 @@ class SessionManager : public ISessionManager { /// Signal all sessions to stop and reject new registrations. /// Clears the cache (destroying idle cached sessions). + /// + /// Cancellation targets every live Session in the process (see LiveSessionRegistry), + /// not just registered ones: a runaway request on a direct-API session never goes + /// through SessionRegistration, yet it is exactly what pins the model loaded. void CancelAll(); /// Block until all sessions have deregistered, with timeout. diff --git a/sdk_v2/cpp/test/internal_api/chat/chat_session_test.cc b/sdk_v2/cpp/test/internal_api/chat/chat_session_test.cc index c6511c769..b0bc37b75 100644 --- a/sdk_v2/cpp/test/internal_api/chat/chat_session_test.cc +++ b/sdk_v2/cpp/test/internal_api/chat/chat_session_test.cc @@ -19,8 +19,11 @@ #include +#include +#include #include #include +#include #include using namespace fl; @@ -338,3 +341,131 @@ TEST_F(ChatSessionTest, SearchOptionsFromEmptyParameters) { EXPECT_FALSE(opts.temperature.has_value()); EXPECT_FALSE(opts.max_output_tokens.has_value()); } + +// =========================================================================== +// Cancellation and deadlines +// +// Regression coverage for the defect where a non-streaming generation could not be +// cancelled or timed out. Because Request::canceled was only polled by the streaming +// loop, a runaway generation pinned the session refcount forever, so Model unload +// failed with "N session(s) still using it" and manager teardown blew its drain +// deadline. +// =========================================================================== + +TEST_F(ChatSessionTest, NonStreamingRequestHonorsTimeout) { + ChatSession session(GetCatalogModel(), GetModel(), *logger_, null_telemetry_); + + Request request; + request.AddOwnedItem(MakeMessage(FOUNDRY_LOCAL_ROLE_USER, "Write an extremely long story.")); + // A large budget the model cannot finish inside the deadline, so the deadline is what + // ends the run rather than a natural stop. + request.options.Add("max_output_tokens", "4096"); + request.SetTimeout(std::chrono::milliseconds(500)); + + Response response; + const auto start = std::chrono::steady_clock::now(); + + EXPECT_THROW(session.ProcessRequest(request, response), fl::Exception); + + const auto elapsed = std::chrono::steady_clock::now() - start; + + EXPECT_TRUE(request.timed_out.load()); + + // Generous ceiling: this asserts the deadline is enforced at all, not its precision. + // The check that matters is that the call returns rather than running to 4096 tokens. + EXPECT_LT(elapsed, std::chrono::seconds(30)); +} + +TEST_F(ChatSessionTest, TimeoutErrorCodeIsTimeout) { + ChatSession session(GetCatalogModel(), GetModel(), *logger_, null_telemetry_); + + Request request; + request.AddOwnedItem(MakeMessage(FOUNDRY_LOCAL_ROLE_USER, "Write an extremely long story.")); + request.options.Add("max_output_tokens", "4096"); + request.SetTimeout(std::chrono::milliseconds(500)); + + Response response; + + try { + session.ProcessRequest(request, response); + FAIL() << "Expected a timeout"; + } catch (const fl::Exception& ex) { + // Distinguishable from a user-initiated cancel so callers can tell "ran out of time" + // from "model stopped early". + EXPECT_EQ(ex.code(), FOUNDRY_LOCAL_ERROR_TIMEOUT); + } +} + +TEST_F(ChatSessionTest, TimeoutIsRearmedPerRequest) { + ChatSession session(GetCatalogModel(), GetModel(), *logger_, null_telemetry_); + + Request request; + request.AddOwnedItem(MakeMessage(FOUNDRY_LOCAL_ROLE_USER, "Write an extremely long story.")); + request.options.Add("max_output_tokens", "4096"); + request.SetTimeout(std::chrono::milliseconds(500)); + + Response response; + EXPECT_THROW(session.ProcessRequest(request, response), fl::Exception); + + // Reusing the object must not leave it permanently in a timed-out state. + Request request2; + request2.AddOwnedItem(MakeMessage(FOUNDRY_LOCAL_ROLE_USER, "What is 2+2? Answer with just the number.")); + request2.options.Add("max_output_tokens", "16"); + request2.options.Add("temperature", "0"); + + Response response2; + EXPECT_NO_THROW(session.ProcessRequest(request2, response2)); + EXPECT_FALSE(request2.timed_out.load()); +} + +TEST_F(ChatSessionTest, CancelStopsInFlightNonStreamingRequest) { + ChatSession session(GetCatalogModel(), GetModel(), *logger_, null_telemetry_); + + Request request; + request.AddOwnedItem(MakeMessage(FOUNDRY_LOCAL_ROLE_USER, "Write an extremely long story.")); + request.options.Add("max_output_tokens", "4096"); + + Response response; + std::atomic finished{false}; + + // Cancel from another thread while ProcessRequest is blocked, which is the whole point: + // the caller of the non-streaming API has no other opportunity to intervene. + std::thread canceller([&]() { + while (!finished.load()) { + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + session.Cancel(); + } + }); + + const auto start = std::chrono::steady_clock::now(); + session.ProcessRequest(request, response); + const auto elapsed = std::chrono::steady_clock::now() - start; + + finished.store(true); + canceller.join(); + + EXPECT_LT(elapsed, std::chrono::seconds(60)); + EXPECT_EQ(response.finish_reason, FOUNDRY_LOCAL_FINISH_NONE); +} + +TEST_F(ChatSessionTest, CancelBeforeRequestRejectsImmediately) { + ChatSession session(GetCatalogModel(), GetModel(), *logger_, null_telemetry_); + session.Cancel(); + + Request request; + request.AddOwnedItem(MakeMessage(FOUNDRY_LOCAL_ROLE_USER, "Hello")); + request.options.Add("max_output_tokens", "16"); + + Response response; + + // A generation started after Cancel() must not run unbounded — otherwise a cancel that + // races with request submission is silently lost. + EXPECT_THROW(session.ProcessRequest(request, response), fl::Exception); +} + +TEST_F(ChatSessionTest, CancelIsIdempotentAndSafeWhenIdle) { + ChatSession session(GetCatalogModel(), GetModel(), *logger_, null_telemetry_); + + EXPECT_NO_THROW(session.Cancel()); + EXPECT_NO_THROW(session.Cancel()); +} diff --git a/sdk_v2/cs/src/Detail/FoundryLocalApi.cs b/sdk_v2/cs/src/Detail/FoundryLocalApi.cs index 32c508a73..df2913166 100644 --- a/sdk_v2/cs/src/Detail/FoundryLocalApi.cs +++ b/sdk_v2/cs/src/Detail/FoundryLocalApi.cs @@ -109,6 +109,11 @@ internal static void CheckStatus(IntPtr status) throw new OperationCanceledException(msg); } + if (code == FlErrorCode.Timeout) + { + throw new TimeoutException(msg); + } + throw new Microsoft.AI.Foundry.Local.FoundryLocalException(msg); } @@ -772,6 +777,15 @@ internal IntPtr ProcessRequest(IntPtr requestPtr) return responsePtr; } + /// + /// Interrupt any request currently running on this session. Thread-safe; intended to be + /// called from a thread other than the one blocked in ProcessRequest. + /// + internal void Cancel() + { + Api.CheckStatus(Api.Inference.SessionCancel(Ptr)); + } + /// Get the number of completed turns in the session. public ulong TurnCount => (ulong)Api.Inference.SessionGetTurnCount(Ptr); diff --git a/sdk_v2/cs/src/Detail/NativeMethods.cs b/sdk_v2/cs/src/Detail/NativeMethods.cs index fa8e920d7..b84986ca8 100644 --- a/sdk_v2/cs/src/Detail/NativeMethods.cs +++ b/sdk_v2/cs/src/Detail/NativeMethods.cs @@ -70,6 +70,7 @@ public enum FlErrorCode InvalidUsage = 4, OperationCancelled = 5, Network = 6, + Timeout = 7, } public enum FlLogLevel @@ -669,6 +670,12 @@ public delegate IntPtr FlInference_SessionSetStreamingCallbackDelegate(IntPtr se [UnmanagedFunctionPointer(CallingConvention.Winapi)] public delegate IntPtr FlInference_SessionUndoTurnsDelegate(IntPtr session, UIntPtr count); +[UnmanagedFunctionPointer(CallingConvention.Winapi)] +public delegate IntPtr FlInference_RequestSetTimeoutMsDelegate(IntPtr request, ulong timeoutMs); + +[UnmanagedFunctionPointer(CallingConvention.Winapi)] +public delegate IntPtr FlInference_SessionCancelDelegate(IntPtr session); + // --- Configuration API (flConfigurationApi) delegates --- [UnmanagedFunctionPointer(CallingConvention.Winapi)] @@ -921,6 +928,11 @@ public struct FlInferenceApi public FlInference_SessionRemoveToolDefinitionDelegate SessionRemoveToolDefinition; public FlInference_SessionGetTurnCountDelegate SessionGetTurnCount; public FlInference_SessionUndoTurnsDelegate SessionUndoTurns; + + // Cancellation / deadlines. Appended to match the native vtable order — these fields + // must stay last, since layout is sequential and positional. + public FlInference_RequestSetTimeoutMsDelegate RequestSetTimeoutMs; + public FlInference_SessionCancelDelegate SessionCancel; } /// Configuration API table. diff --git a/sdk_v2/cs/src/Request.cs b/sdk_v2/cs/src/Request.cs index 7761f0e61..8db4ee645 100644 --- a/sdk_v2/cs/src/Request.cs +++ b/sdk_v2/cs/src/Request.cs @@ -97,6 +97,24 @@ public void Cancel() Api.CheckStatus(Api.Inference.RequestCancel(Ptr)); } + /// + /// Set a wall-clock deadline for this request. Pass (or a + /// negative value) to disable. + /// + /// + /// The deadline covers the whole ProcessRequest call, including prefill, and applies to + /// streaming and non-streaming generation alike. On expiry the run is interrupted + /// mid-compute and ProcessRequest throws a timeout error, so a non-terminating model + /// cannot pin the session and block model unload. The deadline is re-armed on each + /// ProcessRequest call, so a Request may be reused. + /// + public Request SetTimeout(TimeSpan timeout) + { + ulong timeoutMs = timeout > TimeSpan.Zero ? (ulong)timeout.TotalMilliseconds : 0UL; + Api.CheckStatus(Api.Inference.RequestSetTimeoutMs(Ptr, timeoutMs)); + return this; + } + public void Dispose() { if (!_disposed && Ptr != IntPtr.Zero) diff --git a/sdk_v2/cs/src/Session.cs b/sdk_v2/cs/src/Session.cs index 111ae4855..9508546f4 100644 --- a/sdk_v2/cs/src/Session.cs +++ b/sdk_v2/cs/src/Session.cs @@ -130,17 +130,45 @@ public Session SetStreaming(bool enabled) /// /// Process a request and return the complete response. /// + /// + /// genuinely interrupts an in-flight generation: it is registered to + /// cancel the native request, which stops inferencing mid-compute. Without that + /// registration a token could only prevent the work from starting, and a non-terminating + /// generation would keep the session's reference to the model alive indefinitely. + /// public async Task ProcessRequestAsync(Request request, CancellationToken ct = default) { ThrowIfDisposed(); return await Task.Run(() => { + // Dispose the registration before returning so the token cannot cancel a request + // that has already completed and may be reused for a subsequent call. + using var registration = ct.Register(static state => + { + try { ((Request)state!).Cancel(); } catch { } + }, request); + var responsePtr = _session.ProcessRequest(request.Ptr); return new Response(responsePtr); }, ct).ConfigureAwait(false); } + /// + /// Interrupt any request currently running on this session. + /// + /// + /// Thread-safe and intended to be called from a thread other than the one awaiting + /// . Interrupts inferencing mid-compute rather than only + /// between tokens, so the in-flight call returns promptly and releases its reference to the + /// model. Idempotent and safe to call when no request is running. + /// + public void Cancel() + { + ThrowIfDisposed(); + _session.Cancel(); + } + /// /// Process a request with streaming. Returns a whose async /// iterator yields s as they are produced and whose @@ -276,6 +304,12 @@ protected virtual void Dispose(bool disposing) // a use-after-free when Dispose() races with an in-flight ProcessStreamingRequestAsync. try { _activeStreamingCts?.Cancel(); } catch { } + // 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 { } + var streamingTask = _activeStreamingTask; if (streamingTask != null) { diff --git a/sdk_v2/js/native/src/request.cc b/sdk_v2/js/native/src/request.cc index ae2adb436..7959629c3 100644 --- a/sdk_v2/js/native/src/request.cc +++ b/sdk_v2/js/native/src/request.cc @@ -10,6 +10,8 @@ #include +#include +#include #include #include #include @@ -22,6 +24,7 @@ Napi::Function Request::Init(Napi::Env env) { InstanceMethod("addItem", &Request::AddItem), InstanceMethod("setOptions", &Request::SetOptions), InstanceMethod("cancel", &Request::Cancel), + InstanceMethod("setTimeout", &Request::SetTimeout), InstanceMethod("getItemCount", &Request::GetItemCount), InstanceMethod("getItem", &Request::GetItem), }); @@ -95,6 +98,22 @@ Napi::Value Request::Cancel(const Napi::CallbackInfo& info) { }); } +Napi::Value Request::SetTimeout(const Napi::CallbackInfo& info) { + Napi::Env env = info.Env(); + if (impl_ == nullptr) { + return env.Undefined(); + } + if (info.Length() < 1 || !info[0].IsNumber()) { + Napi::TypeError::New(env, "setTimeout(timeoutMs: number)").ThrowAsJavaScriptException(); + return env.Undefined(); + } + const double ms = info[0].As().DoubleValue(); + return CallChecked(env, [&]() -> Napi::Value { + impl_->SetTimeout(std::chrono::milliseconds(ms > 0 ? static_cast(ms) : 0)); + return env.Undefined(); + }); +} + Napi::Value Request::GetItemCount(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (impl_ == nullptr) { diff --git a/sdk_v2/js/native/src/request.h b/sdk_v2/js/native/src/request.h index 7c9f7eebd..7c73445f7 100644 --- a/sdk_v2/js/native/src/request.h +++ b/sdk_v2/js/native/src/request.h @@ -31,6 +31,7 @@ class Request : public Napi::ObjectWrap { Napi::Value AddItem(const Napi::CallbackInfo& info); Napi::Value SetOptions(const Napi::CallbackInfo& info); Napi::Value Cancel(const Napi::CallbackInfo& info); + Napi::Value SetTimeout(const Napi::CallbackInfo& info); Napi::Value GetItemCount(const Napi::CallbackInfo& info); Napi::Value GetItem(const Napi::CallbackInfo& info); diff --git a/sdk_v2/js/native/src/session.cc b/sdk_v2/js/native/src/session.cc index 0d1247463..52bfca079 100644 --- a/sdk_v2/js/native/src/session.cc +++ b/sdk_v2/js/native/src/session.cc @@ -298,6 +298,7 @@ Napi::Function ChatSession::Init(Napi::Env env) { InstanceMethod("removeToolDefinition", &ChatSession::RemoveToolDefinition), InstanceMethod("turnCount", &ChatSession::TurnCount), InstanceMethod("undoTurns", &ChatSession::UndoTurns), + InstanceMethod("cancel", &ChatSession::Cancel), InstanceMethod("dispose", &ChatSession::Dispose), InstanceMethod("isDisposed", &ChatSession::IsDisposed), }); @@ -342,6 +343,19 @@ bool ChatSession::ThrowIfDisposed(Napi::Env env) { return false; } +// Cancel is deliberately synchronous: it must be callable from the JS thread while a +// processRequest promise is still pending on a worker thread. Session::Cancel() only +// signals the engine and returns promptly, so it does not block the event loop. +template +Napi::Value CancelOn(Napi::Env env, SessT* sess) { + try { + sess->Cancel(); + } catch (const std::exception& ex) { + Napi::Error::New(env, ex.what()).ThrowAsJavaScriptException(); + } + return env.Undefined(); +} + Napi::Value ChatSession::ProcessRequest(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (ThrowIfDisposed(env)) return env.Undefined(); @@ -436,6 +450,12 @@ Napi::Value ChatSession::UndoTurns(const Napi::CallbackInfo& info) { }); } +Napi::Value ChatSession::Cancel(const Napi::CallbackInfo& info) { + Napi::Env env = info.Env(); + if (ThrowIfDisposed(env)) return env.Undefined(); + return CancelOn(env, impl_.get()); +} + Napi::Value ChatSession::Dispose(const Napi::CallbackInfo& info) { impl_.reset(); return info.Env().Undefined(); @@ -454,6 +474,7 @@ Napi::Function EmbeddingsSession::Init(Napi::Env env) { { InstanceMethod("processRequest", &EmbeddingsSession::ProcessRequest), InstanceMethod("setOptions", &EmbeddingsSession::SetOptions), + InstanceMethod("cancel", &EmbeddingsSession::Cancel), InstanceMethod("dispose", &EmbeddingsSession::Dispose), InstanceMethod("isDisposed", &EmbeddingsSession::IsDisposed), }); @@ -525,6 +546,12 @@ Napi::Value EmbeddingsSession::SetOptions(const Napi::CallbackInfo& info) { }); } +Napi::Value EmbeddingsSession::Cancel(const Napi::CallbackInfo& info) { + Napi::Env env = info.Env(); + if (ThrowIfDisposed(env)) return env.Undefined(); + return CancelOn(env, impl_.get()); +} + Napi::Value EmbeddingsSession::Dispose(const Napi::CallbackInfo& info) { impl_.reset(); return info.Env().Undefined(); @@ -548,6 +575,7 @@ Napi::Function AudioSession::Init(Napi::Env env) { InstanceMethod("processRequest", &AudioSession::ProcessRequest), InstanceMethod("processStreamingRequest", &AudioSession::ProcessStreamingRequest), InstanceMethod("setOptions", &AudioSession::SetOptions), + InstanceMethod("cancel", &AudioSession::Cancel), InstanceMethod("dispose", &AudioSession::Dispose), InstanceMethod("isDisposed", &AudioSession::IsDisposed), }); @@ -626,6 +654,12 @@ Napi::Value AudioSession::SetOptions(const Napi::CallbackInfo& info) { }); } +Napi::Value AudioSession::Cancel(const Napi::CallbackInfo& info) { + Napi::Env env = info.Env(); + if (ThrowIfDisposed(env)) return env.Undefined(); + return CancelOn(env, impl_.get()); +} + Napi::Value AudioSession::Dispose(const Napi::CallbackInfo& info) { impl_.reset(); return info.Env().Undefined(); diff --git a/sdk_v2/js/native/src/session.h b/sdk_v2/js/native/src/session.h index 2b1db7b86..3a72dc7b9 100644 --- a/sdk_v2/js/native/src/session.h +++ b/sdk_v2/js/native/src/session.h @@ -46,6 +46,7 @@ class ChatSession : public Napi::ObjectWrap { Napi::Value RemoveToolDefinition(const Napi::CallbackInfo& info); Napi::Value TurnCount(const Napi::CallbackInfo& info); Napi::Value UndoTurns(const Napi::CallbackInfo& info); + Napi::Value Cancel(const Napi::CallbackInfo& info); Napi::Value Dispose(const Napi::CallbackInfo& info); Napi::Value IsDisposed(const Napi::CallbackInfo& info); @@ -77,6 +78,7 @@ class EmbeddingsSession : public Napi::ObjectWrap { private: Napi::Value ProcessRequest(const Napi::CallbackInfo& info); Napi::Value SetOptions(const Napi::CallbackInfo& info); + Napi::Value Cancel(const Napi::CallbackInfo& info); Napi::Value Dispose(const Napi::CallbackInfo& info); Napi::Value IsDisposed(const Napi::CallbackInfo& info); @@ -106,6 +108,7 @@ class AudioSession : public Napi::ObjectWrap { Napi::Value ProcessRequest(const Napi::CallbackInfo& info); Napi::Value ProcessStreamingRequest(const Napi::CallbackInfo& info); Napi::Value SetOptions(const Napi::CallbackInfo& info); + Napi::Value Cancel(const Napi::CallbackInfo& info); Napi::Value Dispose(const Napi::CallbackInfo& info); Napi::Value IsDisposed(const Napi::CallbackInfo& info); diff --git a/sdk_v2/js/src/detail/errors.ts b/sdk_v2/js/src/detail/errors.ts index 751dd8c5e..4bbbe4f56 100644 --- a/sdk_v2/js/src/detail/errors.ts +++ b/sdk_v2/js/src/detail/errors.ts @@ -18,6 +18,7 @@ export const FlErrorCode = Object.freeze({ InvalidUsage: 4, OperationCancelled: 5, Network: 6, + Timeout: 7, } as const); export type FlErrorCode = (typeof FlErrorCode)[keyof typeof FlErrorCode]; diff --git a/sdk_v2/js/src/detail/native.ts b/sdk_v2/js/src/detail/native.ts index b47f35f55..3d1c32d9d 100644 --- a/sdk_v2/js/src/detail/native.ts +++ b/sdk_v2/js/src/detail/native.ts @@ -148,6 +148,7 @@ export interface NativeRequest { addItem(item: unknown): void; setOptions(options: NativeRequestOptions): void; cancel(): void; + setTimeout(timeoutMs: number): void; getItemCount(): number; getItem(index: number): unknown; } @@ -169,6 +170,7 @@ export interface NativeSession { processRequest(request: NativeRequest): Promise; processStreamingRequest(request: NativeRequest, onItem: (item: unknown) => void): Promise; setOptions(options: NativeRequestOptions): void; + cancel(): void; dispose(): void; isDisposed(): boolean; } diff --git a/sdk_v2/js/src/request.ts b/sdk_v2/js/src/request.ts index 840812e65..b71e6f7e3 100644 --- a/sdk_v2/js/src/request.ts +++ b/sdk_v2/js/src/request.ts @@ -81,6 +81,23 @@ export class Request { cancel(): void { this.#native.cancel(); } + + /** + * Set a wall-clock deadline for this request, in milliseconds. Pass `0` (or a + * negative value) to disable. + * + * The deadline covers the whole `Session.processRequest()` call, including prefill, + * and applies to streaming and non-streaming generation alike. On expiry the run is + * interrupted mid-compute and the promise rejects with a `FoundryLocalError` whose + * `code === FlErrorCode.Timeout`, so a non-terminating model cannot pin the session + * and block model unload. + * + * The deadline is re-armed on each `processRequest()` call, so a request may be reused. + */ + setTimeout(timeoutMs: number): this { + this.#native.setTimeout(timeoutMs > 0 ? timeoutMs : 0); + return this; + } } /** @internal — used by `Session.processRequest()` to forward the underlying handle. */ diff --git a/sdk_v2/js/src/session.ts b/sdk_v2/js/src/session.ts index 2504fa383..f9cc112c2 100644 --- a/sdk_v2/js/src/session.ts +++ b/sdk_v2/js/src/session.ts @@ -257,10 +257,48 @@ export abstract class Session { * Rejects with a `FoundryLocalError` on native failure. Calling * `request.cancel()` from another async context causes this promise to * reject with `code === FlErrorCode.OperationCancelled`. + * + * Cancellation: pass `{ signal }` to abort a generation that is already running. + * Aborting interrupts inferencing mid-compute, not merely between tokens, so a + * non-terminating generation cannot keep this session's reference to the model alive. */ - async processRequest(request: Request): Promise { + async processRequest(request: Request, options?: StreamOptions): Promise { const nativeReq = unwrapNativeRequest(request); - return (await this.native.processRequest(nativeReq)) as Response; + const signal = options?.signal; + + if (signal?.aborted) { + request.cancel(); + } + + const onAbort = (): void => { + try { + request.cancel(); + } catch { + // The request may already have completed; cancelling it then is a no-op. + } + }; + + signal?.addEventListener("abort", onAbort, { once: true }); + + try { + return (await this.native.processRequest(nativeReq)) as Response; + } finally { + // Detach before returning so the signal cannot cancel a subsequent reuse of + // this request, and so an long-lived signal does not retain it. + signal?.removeEventListener("abort", onAbort); + } + } + + /** + * Interrupt any request currently running on this session. + * + * Unlike `request.cancel()`, this does not require a handle on the in-flight request, + * which makes it the right tool for teardown. It interrupts inferencing mid-compute so + * the pending `processRequest()` promise settles promptly and releases the session's + * reference to the model. Idempotent and safe to call when nothing is running. + */ + cancel(): void { + this.native.cancel(); } /** @@ -299,6 +337,17 @@ export abstract class Session { * with `FoundryLocalError` / `code === FlErrorCode.InvalidUsage`. */ dispose(): void { + // Interrupt anything still running first. Releasing the native session while a worker + // thread is inside processRequest is a use-after-free, and a non-terminating + // generation would otherwise keep the model loaded past teardown. + if (!this.native.isDisposed()) { + try { + this.native.cancel(); + } catch { + // Best-effort: tear down regardless. + } + } + this.native.dispose(); } diff --git a/sdk_v2/python/src/foundry_local_sdk/_native/build_cffi.py b/sdk_v2/python/src/foundry_local_sdk/_native/build_cffi.py index c48ec92ff..94190a0ab 100644 --- a/sdk_v2/python/src/foundry_local_sdk/_native/build_cffi.py +++ b/sdk_v2/python/src/foundry_local_sdk/_native/build_cffi.py @@ -61,6 +61,7 @@ FOUNDRY_LOCAL_ERROR_INVALID_USAGE = 4, FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED = 5, FOUNDRY_LOCAL_ERROR_NETWORK = 6, + FOUNDRY_LOCAL_ERROR_TIMEOUT = 7, } flErrorCode; typedef enum flLogLevel { @@ -456,6 +457,8 @@ flStatusPtr (*Session_RemoveToolDefinition)(flSession* session, const char* tool_name, bool* out_removed); size_t (*Session_GetTurnCount)(const flSession* session); flStatusPtr (*Session_UndoTurns)(flSession* session, size_t count); + flStatusPtr (*Request_SetTimeoutMs)(flRequest* request, uint64_t timeout_ms); + flStatusPtr (*Session_Cancel)(flSession* session); } flInferenceApi; /* ----------------------------------------------------------------------- diff --git a/sdk_v2/python/src/foundry_local_sdk/request.py b/sdk_v2/python/src/foundry_local_sdk/request.py index fe6879e81..9da4f9491 100644 --- a/sdk_v2/python/src/foundry_local_sdk/request.py +++ b/sdk_v2/python/src/foundry_local_sdk/request.py @@ -109,6 +109,27 @@ def cancel(self) -> None: api.check_status(api.inference.Request_Cancel(self._ptr)) + def set_timeout(self, timeout: "float | None") -> "Request": + """Set a wall-clock deadline for this request, in seconds. + + The deadline covers the entire ``process_request`` call, including prefill, and + applies to streaming and non-streaming generation alike. On expiry the run is + interrupted mid-compute and ``process_request`` raises + :class:`~foundry_local_sdk.exceptions.FoundryLocalError` with a timeout code. + + The deadline is re-armed on each ``process_request`` call, so a request object may + be reused. + + Args: + timeout: Deadline in seconds. ``None`` or a non-positive value disables it. + """ + self._check_open() + from foundry_local_sdk._native.api import api + + timeout_ms = 0 if timeout is None or timeout <= 0 else int(timeout * 1000) + api.check_status(api.inference.Request_SetTimeoutMs(self._ptr, timeout_ms)) + return self + def _close(self) -> None: if self._closed: return diff --git a/sdk_v2/python/src/foundry_local_sdk/session.py b/sdk_v2/python/src/foundry_local_sdk/session.py index 136609085..b8fa675ba 100644 --- a/sdk_v2/python/src/foundry_local_sdk/session.py +++ b/sdk_v2/python/src/foundry_local_sdk/session.py @@ -409,17 +409,41 @@ def process_streaming_request( self._streaming_in_flight.release() raise - def process_request(self, request: "Request") -> "Response": - """Run the request synchronously and return the complete response.""" + def process_request(self, request: "Request", timeout: "float | None" = None) -> "Response": + """Run the request synchronously and return the complete response. + + Args: + request: The request to run. + timeout: Optional wall-clock deadline in seconds. On expiry the generation is + interrupted mid-compute and a timeout error is raised, so a non-terminating + model cannot pin the session and block model unload. Overrides any deadline + previously set via :meth:`Request.set_timeout`. + """ self._check_open() from foundry_local_sdk._native import ffi from foundry_local_sdk._native.api import api from foundry_local_sdk.response import Response + if timeout is not None: + request.set_timeout(timeout) + out = ffi.new("flResponse**") api.check_status(api.inference.Session_ProcessRequest(self._ptr, request._ptr, out)) return Response(out[0]) + def cancel(self) -> None: + """Interrupt any request currently running on this session. + + Thread-safe and intended to be called from a thread other than the one blocked in + :meth:`process_request`. Interrupts inferencing mid-compute rather than only + between tokens, so the blocked call returns promptly and releases its reference to + the model. Idempotent and safe to call when nothing is running. + """ + self._check_open() + from foundry_local_sdk._native.api import api + + api.check_status(api.inference.Session_Cancel(self._ptr)) + def _close(self) -> None: # Defensive: subclasses (ChatSession, AudioSession, EmbeddingsSession) validate # the model task BEFORE calling super().__init__(), so a validation failure leaves @@ -428,9 +452,17 @@ def _close(self) -> None: if getattr(self, "_closed", True) or getattr(self, "_ptr", None) is None: return - # If a streaming request is in flight, wind it down before Session_Release — - # releasing while the worker is inside Session_ProcessRequest is a native - # use-after-free. + # Wind down anything still running before Session_Release — releasing while a + # thread is inside Session_ProcessRequest is a native use-after-free. + # + # 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 + t = getattr(self, "_stream_thread", None) if t is not None and t.is_alive(): req = getattr(self, "_stream_request", None)