init

Author: mzegla

src/llm/language_model/legacy/servable.cpp

@@ -83,8 +83,12 @@ absl::Status LegacyServable::parseRequest(std::shared_ptr<GenAiServableExecution
 
     if (legacyExecutionContext->apiHandler->isStream()) {
         legacyExecutionContext->lastStreamerCallbackOutput = "";  // initialize with empty string
-        auto callback = [& executionInProgress = legacyExecutionContext->executionInProgress, &mutex = legacyExecutionContext->mutex, &lastStreamerCallbackOutput = legacyExecutionContext->lastStreamerCallbackOutput](std::string text) {
+        auto callback = [& executionInProgress = legacyExecutionContext->executionInProgress, &mutex = legacyExecutionContext->mutex, &lastStreamerCallbackOutput = legacyExecutionContext->lastStreamerCallbackOutput, &clientDisconnected = legacyExecutionContext->clientDisconnected](std::string text) {
             SPDLOG_LOGGER_TRACE(llm_calculator_logger, "Streamer callback executed with text: [{}]", text);
+            if (clientDisconnected.load()) {
+                executionInProgress.notify_one();
+                return ov::genai::StreamingStatus::CANCEL;
+            }
             {
                 std::lock_guard<std::mutex> lock(mutex);
                 lastStreamerCallbackOutput += text;
@@ -130,10 +134,11 @@ absl::Status LegacyServable::scheduleExecution(std::shared_ptr<GenAiServableExec
     std::weak_ptr<LegacyServableExecutionContext> weakContext = legacyExecutionContext;
     legacyExecutionContext->payload.client->registerDisconnectionCallback([weakContext]() {
         if (auto context = weakContext.lock()) {
-            context->clientDisconnected = true;
+            context->signalDisconnection();
         }
     });
     if (legacyExecutionContext->payload.client->isDisconnected()) {
+        legacyExecutionContext->signalDisconnection();
         return absl::CancelledError();
     }
     properties->legacyExecutor->addRequest(legacyExecutionContext);

src/llm/language_model/legacy/servable.hpp

@@ -37,6 +37,11 @@ struct LegacyServableExecutionContext : public GenAiServableExecutionContext {
 
     // Disconnection handling
     std::atomic<bool> clientDisconnected{false};
+
+    void signalDisconnection() {
+        clientDisconnected = true;
+        executionInProgress.notify_all();
+    }
 };
 
 struct LegacyServableProperties : public GenAiServableProperties {

src/llm/visual_language_model/legacy/servable.cpp

@@ -92,8 +92,12 @@ absl::Status VisualLanguageModelLegacyServable::parseRequest(std::shared_ptr<Gen
 
     if (legacyExecutionContext->apiHandler->isStream()) {
         legacyExecutionContext->lastStreamerCallbackOutput = "";  // initialize with empty string
-        auto callback = [& executionInProgress = legacyExecutionContext->executionInProgress, &mutex = legacyExecutionContext->mutex, &lastStreamerCallbackOutput = legacyExecutionContext->lastStreamerCallbackOutput](std::string text) {
+        auto callback = [& executionInProgress = legacyExecutionContext->executionInProgress, &mutex = legacyExecutionContext->mutex, &lastStreamerCallbackOutput = legacyExecutionContext->lastStreamerCallbackOutput, &clientDisconnected = legacyExecutionContext->clientDisconnected](std::string text) {
             SPDLOG_LOGGER_TRACE(llm_calculator_logger, "Streamer callback executed with text: [{}]", text);
+            if (clientDisconnected.load()) {
+                executionInProgress.notify_one();
+                return ov::genai::StreamingStatus::CANCEL;
+            }
             {
                 std::lock_guard<std::mutex> lock(mutex);
                 lastStreamerCallbackOutput += text;
@@ -123,10 +127,11 @@ absl::Status VisualLanguageModelLegacyServable::scheduleExecution(std::shared_pt
     std::weak_ptr<VisualLanguageModelLegacyServableExecutionContext> weakContext = legacyExecutionContext;
     legacyExecutionContext->payload.client->registerDisconnectionCallback([weakContext]() {
         if (auto context = weakContext.lock()) {
-            context->clientDisconnected = true;
+            context->signalDisconnection();
         }
     });
     if (legacyExecutionContext->payload.client->isDisconnected()) {
+        legacyExecutionContext->signalDisconnection();
         return absl::CancelledError();
     }
     properties->legacyExecutor->addRequest(legacyExecutionContext);

src/llm/visual_language_model/legacy/servable.hpp

@@ -40,6 +40,11 @@ struct VisualLanguageModelLegacyServableExecutionContext : public GenAiServableE
 
     // Disconnection handling
     std::atomic<bool> clientDisconnected{false};
+
+    void signalDisconnection() {
+        clientDisconnected = true;
+        executionInProgress.notify_all();
+    }
 };
 
 struct VisualLanguageModelLegacyServableProperties : public GenAiServableProperties {