Debug registration

router/src/infer.rs

@@ -408,9 +408,17 @@ impl Infer {
         let mut result_start = None;
         let mut result_queued = None;
 
+        tracing::info!("Waiting for response");
+
+        let mut id = None;
+
         // Iterate on stream
         while let Some(response) = stream.next().await {
             match response? {
+                InferStreamResponse::Register { id_val } => {
+                    id = Some(id_val);
+                    tracing::info!("Register response id={id:?}");
+                }
                 // Add prefill tokens
                 InferStreamResponse::Prefill {
                     tokens,
@@ -428,9 +436,13 @@ impl Infer {
                             .collect();
                     }
                     result_prefill_length = tokens_length;
+                    tracing::info!("Prefill response id={id:?}");
                 }
                 // Push last token
-                InferStreamResponse::Token(token) => result_tokens.push(token),
+                InferStreamResponse::Token(token) => {
+                    tracing::info!("Token response id={id:?}");
+                    result_tokens.push(token)
+                }
                 // Final message
                 // Set return values
                 InferStreamResponse::End {
@@ -439,6 +451,7 @@ impl Infer {
                     start,
                     queued,
                 } => {
+                    tracing::info!("End response id={id:?}");
                     result_tokens.push(token);
                     result_generated_text = Some(generated_text);
                     result_start = Some(start);
@@ -455,6 +468,8 @@ impl Infer {
             }
         }
 
+        tracing::info!("Finished response id={id:?}");
+
         // Check that we received a `InferStreamResponse::End` message
         if let (Some(generated_text), Some(queued), Some(start)) =
             (result_generated_text, result_queued, result_start)
@@ -564,6 +579,9 @@ impl Infer {
         let mut stream = UnboundedReceiverStream::new(response_rx);
         while let Some(response) = stream.next().await {
             match response? {
+                InferStreamResponse::Register { .. } => {
+                    tracing::error!("Received a Register message in embed. This is a bug.");
+                }
                 // Add prefill tokens
                 InferStreamResponse::Prefill { .. } => {
                     tracing::error!("Received a Prefill message in embed. This is a bug.");
@@ -667,6 +685,9 @@ impl Infer {
         let mut stream = UnboundedReceiverStream::new(response_rx);
         while let Some(response) = stream.next().await {
             match response? {
+                InferStreamResponse::Register { .. } => {
+                    tracing::error!("Received a Register message in classify. This is a bug.");
+                }
                 // Add prefill tokens
                 InferStreamResponse::Prefill { .. } => {
                     tracing::error!("Received a Prefill message in classify. This is a bug.");
@@ -1350,7 +1371,8 @@ fn send_responses(
 ) -> Result<bool, Box<SendError<Result<InferStreamResponse, InferError>>>> {
     // Return directly if the channel is closed
     if entry.response_tx.is_closed() {
-        tracing::error!("Entry response channel closed.");
+        let id = generation.request_id;
+        tracing::error!("Entry id={id:?} response channel closed.");
         metrics::increment_counter!("lorax_request_failure", "err" => "dropped");
         return Ok(true);
     }
@@ -1497,6 +1519,9 @@ fn send_errors(error: ClientError, entries: &mut IntMap<u64, Entry>) {
 #[derive(Debug)]
 pub(crate) enum InferStreamResponse {
     // Optional first message
+    Register {
+        id_val: u64,
+    },
     Prefill {
         tokens: Option<PrefillTokens>,
         tokens_length: u32,

router/src/queue.rs

@@ -9,7 +9,7 @@ use std::{
 use tokio::{sync::Notify, time::Instant};
 use tracing::info_span;
 
-use crate::{adapter::Adapter, batch::Entry};
+use crate::{adapter::Adapter, batch::Entry, infer::InferStreamResponse};
 
 #[derive(Debug, PartialEq)]
 pub(crate) enum AdapterStatus {
@@ -74,6 +74,11 @@ impl QueueState {
         let queue_span = info_span!(parent: &entry.span, "queued");
         entry.temp_span = Some(queue_span);
 
+        entry
+            .response_tx
+            .send(Ok(InferStreamResponse::Register { id_val: entry_id }))
+            .unwrap();
+
         // Push entry in the queue
         self.entries.push_back((entry_id, entry));
     }
@@ -214,9 +219,12 @@ impl AdapterQueuesState {
 
         // ensure that append completes before sending batcher message
         let queue = self.queue_map.get_mut(&adapter).unwrap();
+        let id = self.next_id;
         queue.append(self.next_id, entry);
         self.next_id += 1;
 
+        tracing::info!("append entry id={:?} adapter={:?}", id, adapter.index());
+
         return download;
     }
 

router/src/server.rs

@@ -1062,6 +1062,11 @@ async fn generate_stream_with_callback(
                             match response {
                                 Ok(response) => {
                                     match response {
+                                        InferStreamResponse::Register {
+                                            ..
+                                        } => {
+                                            // Register is ignored
+                                        }
                                         // Prefill is ignored
                                         InferStreamResponse::Prefill {
                                             tokens_length,