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()) {
+                std::cout << "Notifying last time" << std::endl;
+                request->notify_handle();
+                std::cout << "Dropping" << std::endl;
                 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) {
@@ -149,6 +152,7 @@ class ContinuousBatchingPipeline::Impl {
 
     void step() {
         static ManualTimer step_timer("step()");
+        std::cout << "Step()" << std::endl;
         step_timer.start();
 
         // Pull awaiting requests

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();
 }
@@ -44,7 +48,10 @@ std::vector<GenerationOutput> GenerationHandleImpl::read_all() {
     std::vector<GenerationOutput> results;
     std::unordered_map<uint64_t, GenerationOutput> partial_results;
     // We iterate until generation is running or there are tokens we haven't read yet
+    std::cout << "AAAAAAAAAAAAAAAAAAAAAAAAAA Starting" << std::endl;
+    int i = 0;
     while (get_status() == GenerationStatus::RUNNING || can_read()) {
+        std::cout << "AAAAAAAAAAAAAAAAAAAAAAAAAA Iteration" << i++ << std::endl;
         // For unary case there's only one iteration and we get all results in a single read() call
         std::unordered_map<uint64_t, GenerationOutput> iteration_results = read();
         add_partial_result(partial_results, iteration_results);

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()) {