V1

src/cpp/include/openvino/genai/generation_handle.hpp

@@ -74,12 +74,14 @@ class OPENVINO_GENAI_EXPORTS GenerationHandleImpl {
 
     bool can_read();
 
+    void drop();
+
     GenerationOutputs back();
     // Reads result of a generation for single iteration
     GenerationOutputs read();
     // Reads all generated tokens for all sequences
     std::vector<GenerationOutput> read_all();
 };
 
-using GenerationHandle = std::unique_ptr<GenerationHandleImpl>;
+using GenerationHandle = std::shared_ptr<GenerationHandleImpl>;
 }

src/cpp/src/continuous_batching_pipeline.cpp

@@ -65,6 +65,9 @@ class ContinuousBatchingPipeline::Impl {
         while (requests_iterator != m_requests.end()) {
             const auto& request = *requests_iterator;
             if(request->has_finished() || request->out_of_memory() || request->handle_dropped()) {
+                // Notify the last time even if there will be no results
+                // This causes read_all() to unblock in all situations
+                request->notify_handle();
                 for (const auto& sequence: request->get_sequences()) {
                     m_scheduler->free_sequence(sequence->get_id());
                 }
@@ -136,7 +139,7 @@ class ContinuousBatchingPipeline::Impl {
             std::lock_guard<std::mutex> lock{m_awaiting_requests_mutex};
             m_awaiting_requests.push_back(sequence_group);
         }
-        return std::make_unique<GenerationHandleImpl>(sequence_group->get_generation_stream(), sampling_params);
+        return std::make_shared<GenerationHandleImpl>(sequence_group->get_generation_stream(), sampling_params);
     }
 
     GenerationHandle add_request(uint64_t request_id, const std::string& prompt, ov::genai::GenerationConfig sampling_params) {

src/cpp/src/generation_handle.cpp

@@ -9,7 +9,7 @@
 using namespace ov::genai;
 
 GenerationHandleImpl::~GenerationHandleImpl() {
-    m_generation_stream->drop();
+    drop();
 }
 
 GenerationStatus GenerationHandleImpl::get_status() {
@@ -20,6 +20,10 @@ bool GenerationHandleImpl::can_read() {
     return m_generation_stream->can_read();
 }
 
+void GenerationHandleImpl::drop() {
+    m_generation_stream->drop();
+}
+
 std::unordered_map<uint64_t, GenerationOutput> GenerationHandleImpl::back() {
     return m_generation_stream->back();
 }

src/cpp/src/sequence_group.hpp

@@ -477,7 +477,7 @@ class SequenceGroup {
         }
         // For beam search streaming is not available, so we notify only upon finishing
         if(m_sampling_params.is_beam_search()) {
-            if (has_finished() || out_of_memory()) {
+            if (has_finished() || out_of_memory() || handle_dropped()) {
                 push_outputs();
             }
         } else if (m_sampling_params.is_greedy_decoding() || m_sampling_params.is_multinomial()) {